From 20773fbc36423c9cefc85d20f2220620a514b73f Mon Sep 17 00:00:00 2001 From: Sam Brenner Date: Tue, 24 Sep 2024 09:34:58 -0400 Subject: [PATCH 1/6] tagger --- packages/dd-trace/src/llmobs/constants.js | 26 +++ packages/dd-trace/src/llmobs/tagger.js | 206 ++++++++++++++++++++++ 2 files changed, 232 insertions(+) create mode 100644 packages/dd-trace/src/llmobs/constants.js create mode 100644 packages/dd-trace/src/llmobs/tagger.js diff --git a/packages/dd-trace/src/llmobs/constants.js b/packages/dd-trace/src/llmobs/constants.js new file mode 100644 index 00000000000..ea7d90a9501 --- /dev/null +++ b/packages/dd-trace/src/llmobs/constants.js @@ -0,0 +1,26 @@ +'use strict' + +module.exports = { + SPAN_KIND: '_ml_obs.meta.span.kind', + SESSION_ID: '_ml_obs.session_id', + METADATA: '_ml_obs.meta.metadata', + METRICS: '_ml_obs.metrics', + ML_APP: '_ml_obs.meta.ml_app', + PROPAGATED_PARENT_ID_KEY: '_dd.p.llmobs_parent_id', + PARENT_ID_KEY: '_ml_obs.llmobs_parent_id', + TAGS: '_ml_obs.tags', + NAME: '_ml_obs.name', + TRACE_ID: '_ml_obs.trace_id', + PROPAGATED_TRACE_ID_KEY: '_dd.p.llmobs_trace_id', + + MODEL_NAME: '_ml_obs.meta.model_name', + MODEL_PROVIDER: '_ml_obs.meta.model_provider', + + INPUT_DOCUMENTS: '_ml_obs.meta.input.documents', + INPUT_MESSAGES: '_ml_obs.meta.input.messages', + INPUT_VALUE: '_ml_obs.meta.input.value', + + OUTPUT_DOCUMENTS: '_ml_obs.meta.output.documents', + OUTPUT_MESSAGES: '_ml_obs.meta.output.messages', + OUTPUT_VALUE: '_ml_obs.meta.output.value' +} diff --git a/packages/dd-trace/src/llmobs/tagger.js b/packages/dd-trace/src/llmobs/tagger.js new file mode 100644 index 00000000000..c726da3e2b5 --- /dev/null +++ b/packages/dd-trace/src/llmobs/tagger.js @@ -0,0 +1,206 @@ +'use strict' + +const logger = require('../log') +const { + SPAN_TYPE +} = require('../../../../ext/tags') +const { + MODEL_NAME, + MODEL_PROVIDER, + SESSION_ID, + ML_APP, + SPAN_KIND, + INPUT_VALUE, + OUTPUT_DOCUMENTS, + INPUT_DOCUMENTS, + OUTPUT_VALUE, + METADATA, + METRICS, + PARENT_ID_KEY, + INPUT_MESSAGES, + OUTPUT_MESSAGES, + TAGS, + NAME, + PROPAGATED_PARENT_ID_KEY +} = require('./constants') + +class LLMObsTagger { + constructor (config) { + this._config = config + } + + setLLMObsSpanTags ( + span, + kind, + { modelName, modelProvider, sessionId, mlApp, parentLLMObsSpan } = {}, + name + ) { + if (kind) span.setTag(SPAN_TYPE, 'llm') // only mark it as an llm span if it was a valid kind + if (name) span.setTag(NAME, name) + + span.setTag(SPAN_KIND, kind) + if (modelName) span.setTag(MODEL_NAME, modelName) + if (modelProvider) span.setTag(MODEL_PROVIDER, modelProvider) + + sessionId = sessionId || parentLLMObsSpan?.context()._tags[SESSION_ID] + if (sessionId) span.setTag(SESSION_ID, sessionId) + + if (!mlApp) mlApp = parentLLMObsSpan?.context()._tags[ML_APP] || this._config.llmobs.mlApp + span.setTag(ML_APP, mlApp) + + const parentId = + parentLLMObsSpan?.context().toSpanId() || + span.context()._trace.tags[PROPAGATED_PARENT_ID_KEY] || + 'undefined' + span.setTag(PARENT_ID_KEY, parentId) + } + + tagLLMIO (span, inputData, outputData) { + this._tagMessages(span, inputData, INPUT_MESSAGES) + this._tagMessages(span, outputData, OUTPUT_MESSAGES) + } + + tagEmbeddingIO (span, inputData, outputData) { + this._tagDocuments(span, inputData, INPUT_DOCUMENTS) + this._tagText(span, outputData, OUTPUT_VALUE) + } + + tagRetrievalIO (span, inputData, outputData) { + this._tagText(span, inputData, INPUT_VALUE) + this._tagDocuments(span, outputData, OUTPUT_DOCUMENTS) + } + + tagTextIO (span, inputData, outputData) { + this._tagText(span, inputData, INPUT_VALUE) + this._tagText(span, outputData, OUTPUT_VALUE) + } + + tagMetadata (span, metadata) { + try { + span.setTag(METADATA, JSON.stringify(metadata)) + } catch { + logger.warn('Failed to parse span metadata. Metadata key-value pairs must be JSON serializable.') + } + } + + tagMetrics (span, metrics) { + try { + span.setTag(METRICS, JSON.stringify(metrics)) + } catch { + logger.warn('Failed to parse span metrics. Metrics key-value pairs must be JSON serializable.') + } + } + + tagSpanTags (span, tags) { + try { + const currentTags = span.context()._tags[TAGS] + if (currentTags) { + Object.assign(tags, currentTags) + } + span.setTag(TAGS, JSON.stringify(tags)) + } catch { + logger.warn('Failed to parse span tags. Tag key-value pairs must be JSON serializable.') + } + } + + _tagText (span, data, key) { + if (data) { + if (typeof data === 'string') { + span.setTag(key, data) + } else { + try { + // this will help showcase unfinished promises being passed in as values + span.setTag(key, JSON.stringify(data)) + } catch { + const type = key === INPUT_VALUE ? 'input' : 'output' + logger.warn(`Failed to parse ${type} value, must be JSON serializable.`) + } + } + } + } + + _tagDocuments (span, data, key) { + if (data) { + if (!Array.isArray(data)) { + data = [data] + } + + try { + const documents = data.map(document => { + if (typeof document === 'string') { + return document + } + + const { text, name, id, score } = document + + if (text && typeof text !== 'string') { + logger.warn(`Invalid property found in ${document}: text must be a string.`) + return undefined + } + + if (name && typeof name !== 'string') { + logger.warn(`Invalid property found in ${document}: name must be a string.`) + return undefined + } + + if (id && typeof id !== 'string') { + logger.warn(`Invalid property found in ${document}: id must be a string.`) + return undefined + } + + if (score && typeof score !== 'number') { + logger.warn(`Invalid property found in ${document}: score must be a number.`) + return undefined + } + + return document + }).filter(doc => !!doc) // filter out bad documents? + + span.setTag(key, JSON.stringify(documents)) + } catch { + const type = key === INPUT_DOCUMENTS ? 'input' : 'output' + logger.warn(`Failed to parse ${type} documents.`) + } + } + } + + _tagMessages (span, data, key) { + if (data) { + if (!Array.isArray(data)) { + data = [data] + } + + try { + const messages = data.map(message => { + if (typeof message === 'string') { + return message + } + + const content = message.content || '' + const role = message.role + + if (typeof content !== 'string') { + logger.warn(`Invalid property found in ${message}: content must be a string.`) + return undefined + } + + message.content = content + + if (role && typeof role !== 'string') { + logger.warn(`Invalid property found in ${message}: role must be a string.`) + return undefined + } + + return message + }).filter(msg => !!msg) // filter out bad messages? + + span.setTag(key, JSON.stringify(messages)) + } catch { + const type = key === INPUT_MESSAGES ? 'input' : 'output' + logger.warn(`Failed to parse ${type} messages.`) + } + } + } +} + +module.exports = LLMObsTagger From 1ef14c1e224c437b203fdbbee7980813cb3a612d Mon Sep 17 00:00:00 2001 From: Sam Brenner Date: Tue, 24 Sep 2024 09:35:05 -0400 Subject: [PATCH 2/6] tests --- package.json | 2 + packages/dd-trace/test/llmobs/tagger.spec.js | 430 +++++++++++++++++++ 2 files changed, 432 insertions(+) create mode 100644 packages/dd-trace/test/llmobs/tagger.spec.js diff --git a/package.json b/package.json index 67d4c132a19..ed8295e01ef 100644 --- a/package.json +++ b/package.json @@ -30,6 +30,8 @@ "test:core:ci": "npm run test:core -- --coverage --nyc-arg=--include=\"packages/datadog-core/src/**/*.js\"", "test:lambda": "mocha -r \"packages/dd-trace/test/setup/mocha.js\" \"packages/dd-trace/test/lambda/**/*.spec.js\"", "test:lambda:ci": "nyc --no-clean --include \"packages/dd-trace/src/lambda/**/*.js\" -- npm run test:lambda", + "test:llmobs": "mocha -r \"packages/dd-trace/test/setup/mocha.js\" \"packages/dd-trace/test/llmobs/**/*.spec.js\"", + "test:llmobs:ci": "nyc --no-clean --include \"packages/dd-trace/src/llmobs/**/*.js\" -- npm run test:llmobs", "test:plugins": "mocha -r \"packages/dd-trace/test/setup/mocha.js\" \"packages/datadog-instrumentations/test/@($(echo $PLUGINS)).spec.js\" \"packages/datadog-plugin-@($(echo $PLUGINS))/test/**/*.spec.js\"", "test:plugins:ci": "yarn services && nyc --no-clean --include \"packages/datadog-instrumentations/src/@($(echo $PLUGINS)).js\" --include \"packages/datadog-instrumentations/src/@($(echo $PLUGINS))/**/*.js\" --include \"packages/datadog-plugin-@($(echo $PLUGINS))/src/**/*.js\" -- npm run test:plugins", "test:plugins:upstream": "node ./packages/dd-trace/test/plugins/suite.js", diff --git a/packages/dd-trace/test/llmobs/tagger.spec.js b/packages/dd-trace/test/llmobs/tagger.spec.js new file mode 100644 index 00000000000..16cd607e085 --- /dev/null +++ b/packages/dd-trace/test/llmobs/tagger.spec.js @@ -0,0 +1,430 @@ +'use strict' + +const { expect } = require('chai') +const proxyquire = require('proxyquire') + +function unserializbleObject () { + const obj = {} + obj.obj = obj + return obj +} + +describe('tagger', () => { + let span + let spanContext + let Tagger + let tagger + let logger + let util + + beforeEach(() => { + spanContext = { + _tags: {}, + _trace: { tags: {} } + } + + span = { + context () { return spanContext }, + setTag (k, v) { + this.context()._tags[k] = v + } + } + + util = { + generateTraceId: sinon.stub().returns('0123') + } + + logger = { + warn: sinon.stub() + } + + Tagger = proxyquire('../../src/llmobs/tagger', { + '../log': logger, + './util': util + }) + + tagger = new Tagger({ llmobs: { mlApp: 'my-default-ml-app' } }) + }) + + describe('setLLMObsSpanTags', () => { + it('tags an llm obs span with basic and default properties', () => { + tagger.setLLMObsSpanTags(span, 'workflow') + + expect(span.context()._tags).to.deep.equal({ + 'span.type': 'llm', + '_ml_obs.meta.span.kind': 'workflow', + '_ml_obs.meta.ml_app': 'my-default-ml-app', + '_ml_obs.llmobs_parent_id': 'undefined' // no parent id provided + }) + }) + + it('uses options passed in to set tags', () => { + tagger.setLLMObsSpanTags(span, 'llm', { + modelName: 'my-model', + modelProvider: 'my-provider', + sessionId: 'my-session', + mlApp: 'my-app' + }) + + expect(span.context()._tags).to.deep.equal({ + 'span.type': 'llm', + '_ml_obs.meta.span.kind': 'llm', + '_ml_obs.meta.model_name': 'my-model', + '_ml_obs.meta.model_provider': 'my-provider', + '_ml_obs.session_id': 'my-session', + '_ml_obs.meta.ml_app': 'my-app', + '_ml_obs.llmobs_parent_id': 'undefined' + }) + }) + + it('uses the name if provided', () => { + tagger.setLLMObsSpanTags(span, 'llm', {}, 'my-span-name') + + expect(span.context()._tags).to.deep.equal({ + 'span.type': 'llm', + '_ml_obs.meta.span.kind': 'llm', + '_ml_obs.meta.ml_app': 'my-default-ml-app', + '_ml_obs.llmobs_parent_id': 'undefined', + '_ml_obs.name': 'my-span-name' + }) + }) + + it('defaults parent id to undefined', () => { + tagger.setLLMObsSpanTags(span, 'llm') + + expect(span.context()._tags).to.deep.equal({ + 'span.type': 'llm', + '_ml_obs.meta.span.kind': 'llm', + '_ml_obs.meta.ml_app': 'my-default-ml-app', + '_ml_obs.llmobs_parent_id': 'undefined' + }) + }) + + it('uses the parent span if provided to populate fields', () => { + const parentSpan = { + context () { + return { + _tags: { + '_ml_obs.meta.ml_app': 'my-ml-app', + '_ml_obs.session_id': 'my-session' + }, + toSpanId () { return '5678' } + } + } + } + tagger.setLLMObsSpanTags(span, 'llm', { parentLLMObsSpan: parentSpan }) + + expect(span.context()._tags).to.deep.equal({ + 'span.type': 'llm', + '_ml_obs.meta.span.kind': 'llm', + '_ml_obs.meta.ml_app': 'my-ml-app', + '_ml_obs.session_id': 'my-session', + '_ml_obs.llmobs_parent_id': '5678' + }) + }) + + it('uses the propagated trace id if provided', () => { + spanContext._trace.tags['_dd.p.llmobs_trace_id'] = '-123' + + tagger.setLLMObsSpanTags(span, 'llm') + + expect(span.context()._tags).to.deep.equal({ + 'span.type': 'llm', + '_ml_obs.meta.span.kind': 'llm', + '_ml_obs.meta.ml_app': 'my-default-ml-app', + '_ml_obs.llmobs_parent_id': 'undefined' + }) + }) + + it('uses the propagated parent id if provided', () => { + spanContext._trace.tags['_dd.p.llmobs_parent_id'] = '-567' + + tagger.setLLMObsSpanTags(span, 'llm') + + expect(span.context()._tags).to.deep.equal({ + 'span.type': 'llm', + '_ml_obs.meta.span.kind': 'llm', + '_ml_obs.meta.ml_app': 'my-default-ml-app', + '_ml_obs.llmobs_parent_id': '-567' + }) + }) + + it('does not set span type if the LLMObs span kind is falsy', () => { + tagger.setLLMObsSpanTags(span, false) + + expect(span.context()._tags['span.type']).to.be.undefined + }) + }) + + describe('tagMetadata', () => { + it('tags a span with metadata', () => { + tagger.tagMetadata(span, { a: 'foo', b: 'bar' }) + expect(span.context()._tags).to.deep.equal({ + '_ml_obs.meta.metadata': '{"a":"foo","b":"bar"}' + }) + }) + + it('logs when metadata is not JSON serializable', () => { + const metadata = unserializbleObject() + tagger.tagMetadata(span, metadata) + expect(logger.warn).to.have.been.calledOnce + }) + }) + + describe('tagMetrics', () => { + it('tags a span with metrics', () => { + tagger.tagMetadata(span, { a: 1, b: 2 }) + expect(span.context()._tags).to.deep.equal({ + '_ml_obs.meta.metadata': '{"a":1,"b":2}' + }) + }) + + it('logs when metrics is not JSON serializable', () => { + const metadata = unserializbleObject() + tagger.tagMetadata(span, metadata) + expect(logger.warn).to.have.been.calledOnce + }) + }) + + describe('tagLLMIO', () => { + beforeEach(() => { + sinon.stub(tagger, '_tagMessages') + }) + + afterEach(() => { + tagger._tagMessages.restore() + }) + + it('tags a span with llm io', () => { + tagger.tagLLMIO(span, { a: 'foo' }, { b: 'bar' }) + expect(tagger._tagMessages).to.have.been.calledTwice + expect(tagger._tagMessages).to.have.been.calledWith(span, { a: 'foo' }, '_ml_obs.meta.input.messages') + expect(tagger._tagMessages).to.have.been.calledWith(span, { b: 'bar' }, '_ml_obs.meta.output.messages') + }) + }) + + describe('tagEmbeddingIO', () => { + beforeEach(() => { + sinon.stub(tagger, '_tagDocuments') + sinon.stub(tagger, '_tagText') + }) + + afterEach(() => { + tagger._tagDocuments.restore() + tagger._tagText.restore() + }) + + it('tags a span with embedding io', () => { + tagger.tagEmbeddingIO(span, { a: 'foo' }, { b: 'bar' }) + expect(tagger._tagDocuments).to.have.been.calledOnce + expect(tagger._tagDocuments).to.have.been.calledWith(span, { a: 'foo' }, '_ml_obs.meta.input.documents') + expect(tagger._tagText).to.have.been.calledOnce + expect(tagger._tagText).to.have.been.calledWith(span, { b: 'bar' }, '_ml_obs.meta.output.value') + }) + }) + + describe('tagRetrievalIO', () => { + beforeEach(() => { + sinon.stub(tagger, '_tagDocuments') + sinon.stub(tagger, '_tagText') + }) + + afterEach(() => { + tagger._tagDocuments.restore() + tagger._tagText.restore() + }) + + it('tags a span with retrieval io', () => { + tagger.tagRetrievalIO(span, { a: 'foo' }, { b: 'bar' }) + expect(tagger._tagDocuments).to.have.been.calledOnce + expect(tagger._tagDocuments).to.have.been.calledWith(span, { b: 'bar' }, '_ml_obs.meta.output.documents') + expect(tagger._tagText).to.have.been.calledOnce + expect(tagger._tagText).to.have.been.calledWith(span, { a: 'foo' }, '_ml_obs.meta.input.value') + }) + }) + + describe('tagTextIO', () => { + beforeEach(() => { + sinon.stub(tagger, '_tagText') + }) + + afterEach(() => { + tagger._tagText.restore() + }) + + it('tags a span with text io', () => { + tagger.tagTextIO(span, { a: 'foo' }, { b: 'bar' }) + expect(tagger._tagText).to.have.been.calledTwice + expect(tagger._tagText).to.have.been.calledWith(span, { a: 'foo' }, '_ml_obs.meta.input.value') + expect(tagger._tagText).to.have.been.calledWith(span, { b: 'bar' }, '_ml_obs.meta.output.value') + }) + }) + + // maybe confirm this one + describe('tagSpanTags', () => {}) + + describe('_tagText', () => { + it('tags a span with text', () => { + tagger._tagText(span, 'my-text', '_ml_obs.meta.input.value') + + expect(span.context()._tags).to.deep.equal({ + '_ml_obs.meta.input.value': 'my-text' + }) + }) + + it('tags a span with an object', () => { + tagger._tagText(span, { a: 1, b: 2 }, '_ml_obs.meta.input.value') + + expect(span.context()._tags).to.deep.equal({ + '_ml_obs.meta.input.value': '{"a":1,"b":2}' + }) + }) + + it('logs when the value is not JSON serializable', () => { + const data = unserializbleObject() + tagger._tagText(span, data, '_ml_obs.meta.input.value') + expect(logger.warn).to.have.been.calledOnce + }) + }) + + describe('_tagDocuments', () => { + it('tags a single document object', () => { + const document = { text: 'my-text', name: 'my-name', id: 'my-id', score: 0.5 } + tagger._tagDocuments(span, document, '_ml_obs.meta.input.documents') + + expect(span.context()._tags).to.deep.equal({ + '_ml_obs.meta.input.documents': '[{"text":"my-text","name":"my-name","id":"my-id","score":0.5}]' + }) + }) + + it('tags a single document string', () => { + const document = 'my-text' + tagger._tagDocuments(span, document, '_ml_obs.meta.input.documents') + + expect(span.context()._tags).to.deep.equal({ + '_ml_obs.meta.input.documents': '["my-text"]' + }) + }) + + it('tags a document with a subset of properties', () => { + const document = { text: 'my-text', score: 0.5 } + tagger._tagDocuments(span, document, '_ml_obs.meta.input.documents') + + expect(span.context()._tags).to.deep.equal({ + '_ml_obs.meta.input.documents': '[{"text":"my-text","score":0.5}]' + }) + }) + + it('tags multiple documents', () => { + const documents = [ + { text: 'my-text', name: 'my-name', id: 'my-id', score: 0.5 }, + { text: 'my-text2', name: 'my-name2', id: 'my-id2', score: 0.6 } + ] + tagger._tagDocuments(span, documents, '_ml_obs.meta.input.documents') + + expect(span.context()._tags).to.deep.equal({ + '_ml_obs.meta.input.documents': '[{"text":"my-text","name":"my-name"' + + ',"id":"my-id","score":0.5},{"text":"my-text2","name":"my-name2","id":"my-id2","score":0.6}]' + }) + }) + + it('does not include malformed documents', () => { + const documents = [ + { text: 'my-text', name: 'my-name', id: 'my-id', score: 0.5 }, + { text: 'my-text2', score: 'not-a-number' } + ] + + tagger._tagDocuments(span, documents, '_ml_obs.meta.input.documents') + + expect(span.context()._tags).to.deep.equal({ + '_ml_obs.meta.input.documents': '[{"text":"my-text","name":"my-name","id":"my-id","score":0.5}]' + }) + }) + + it('logs when a document is not JSON serializable', () => { + const documents = [unserializbleObject()] + tagger._tagDocuments(span, documents, '_ml_obs.meta.input.documents') + expect(logger.warn).to.have.been.calledOnce + }) + }) + + describe('_tagMessages', () => { + it('tags a single message object', () => { + const message = { content: 'my-content', role: 'my-role' } + tagger._tagMessages(span, message, '_ml_obs.meta.input.messages') + + expect(span.context()._tags).to.deep.equal({ + '_ml_obs.meta.input.messages': '[{"content":"my-content","role":"my-role"}]' + }) + }) + + it('tags a single message string', () => { + tagger._tagMessages(span, 'my-message', '_ml_obs.meta.input.messages') + + expect(span.context()._tags).to.deep.equal({ + '_ml_obs.meta.input.messages': '["my-message"]' + }) + }) + + it('tags a message with only content', () => { + const message = { content: 'my-content' } + tagger._tagMessages(span, message, '_ml_obs.meta.input.messages') + + expect(span.context()._tags).to.deep.equal({ + '_ml_obs.meta.input.messages': '[{"content":"my-content"}]' + }) + }) + + it('defaults missing content to empty content', () => { + const message = {} + tagger._tagMessages(span, message, '_ml_obs.meta.input.messages') + + expect(span.context()._tags).to.deep.equal({ + '_ml_obs.meta.input.messages': '[{"content":""}]' + }) + }) + + it('filters out messages with non-string content', () => { + const messages = [ + { content: 'my-content' }, + { role: 'my-role', content: 6 } + ] + tagger._tagMessages(span, messages, '_ml_obs.meta.input.messages') + + expect(span.context()._tags).to.deep.equal({ + '_ml_obs.meta.input.messages': '[{"content":"my-content"}]' + }) + }) + + it('filters out messages with non-string role', () => { + const messages = [ + { content: 'my-content', role: 'my-role' }, + { content: 'my-content2', role: 6 } + ] + tagger._tagMessages(span, messages, '_ml_obs.meta.input.messages') + + expect(span.context()._tags).to.deep.equal({ + '_ml_obs.meta.input.messages': '[{"content":"my-content","role":"my-role"}]' + }) + }) + + it('tags multiple messages', () => { + const messages = [ + { content: 'my-content', role: 'my-role' }, + { content: 'my-content2', role: 'my-role2' } + ] + tagger._tagMessages(span, messages, '_ml_obs.meta.input.messages') + + expect(span.context()._tags).to.deep.equal({ + '_ml_obs.meta.input.messages': '[{"content":"my-content","role":"my-role"},' + + '{"content":"my-content2","role":"my-role2"}]' + }) + }) + + it('logs when a message is not JSON serializable', () => { + const messages = [unserializbleObject()] + tagger._tagMessages(span, messages, '_ml_obs.meta.input.messages') + expect(logger.warn).to.have.been.calledOnce + }) + }) +}) From 88c441e4ddf71043d0f7e3a067f0133c716c1cac Mon Sep 17 00:00:00 2001 From: Sam Brenner Date: Tue, 24 Sep 2024 17:26:40 -0400 Subject: [PATCH 3/6] review comments --- packages/dd-trace/src/llmobs/constants.js | 1 + packages/dd-trace/src/llmobs/tagger.js | 5 +++-- packages/dd-trace/test/llmobs/tagger.spec.js | 2 -- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/packages/dd-trace/src/llmobs/constants.js b/packages/dd-trace/src/llmobs/constants.js index ea7d90a9501..3052056d437 100644 --- a/packages/dd-trace/src/llmobs/constants.js +++ b/packages/dd-trace/src/llmobs/constants.js @@ -12,6 +12,7 @@ module.exports = { NAME: '_ml_obs.name', TRACE_ID: '_ml_obs.trace_id', PROPAGATED_TRACE_ID_KEY: '_dd.p.llmobs_trace_id', + ROOT_PARENT_ID: 'undefined', MODEL_NAME: '_ml_obs.meta.model_name', MODEL_PROVIDER: '_ml_obs.meta.model_provider', diff --git a/packages/dd-trace/src/llmobs/tagger.js b/packages/dd-trace/src/llmobs/tagger.js index c726da3e2b5..b203abe5147 100644 --- a/packages/dd-trace/src/llmobs/tagger.js +++ b/packages/dd-trace/src/llmobs/tagger.js @@ -21,7 +21,8 @@ const { OUTPUT_MESSAGES, TAGS, NAME, - PROPAGATED_PARENT_ID_KEY + PROPAGATED_PARENT_ID_KEY, + ROOT_PARENT_ID } = require('./constants') class LLMObsTagger { @@ -51,7 +52,7 @@ class LLMObsTagger { const parentId = parentLLMObsSpan?.context().toSpanId() || span.context()._trace.tags[PROPAGATED_PARENT_ID_KEY] || - 'undefined' + ROOT_PARENT_ID span.setTag(PARENT_ID_KEY, parentId) } diff --git a/packages/dd-trace/test/llmobs/tagger.spec.js b/packages/dd-trace/test/llmobs/tagger.spec.js index 16cd607e085..d66b5e257f1 100644 --- a/packages/dd-trace/test/llmobs/tagger.spec.js +++ b/packages/dd-trace/test/llmobs/tagger.spec.js @@ -124,8 +124,6 @@ describe('tagger', () => { }) it('uses the propagated trace id if provided', () => { - spanContext._trace.tags['_dd.p.llmobs_trace_id'] = '-123' - tagger.setLLMObsSpanTags(span, 'llm') expect(span.context()._tags).to.deep.equal({ From 1667e6e42b9fd0197076d0e89c01fbd0a034edab Mon Sep 17 00:00:00 2001 From: Sam Brenner Date: Thu, 26 Sep 2024 10:42:24 -0400 Subject: [PATCH 4/6] fix tagger & tests --- packages/dd-trace/src/llmobs/constants.js | 2 +- packages/dd-trace/src/llmobs/tagger.js | 73 +++-- packages/dd-trace/test/llmobs/tagger.spec.js | 317 ++++++++----------- 3 files changed, 175 insertions(+), 217 deletions(-) diff --git a/packages/dd-trace/src/llmobs/constants.js b/packages/dd-trace/src/llmobs/constants.js index ebb4437abda..4cc63a8d01d 100644 --- a/packages/dd-trace/src/llmobs/constants.js +++ b/packages/dd-trace/src/llmobs/constants.js @@ -24,7 +24,7 @@ module.exports = { OUTPUT_DOCUMENTS: '_ml_obs.meta.output.documents', OUTPUT_MESSAGES: '_ml_obs.meta.output.messages', OUTPUT_VALUE: '_ml_obs.meta.output.value', - + EVP_PROXY_AGENT_BASE_PATH: 'evp_proxy/v2', EVP_PROXY_AGENT_ENDPOINT: 'evp_proxy/v2/api/v2/llmobs', EVP_SUBDOMAIN_HEADER_NAME: 'X-Datadog-EVP-Subdomain', diff --git a/packages/dd-trace/src/llmobs/tagger.js b/packages/dd-trace/src/llmobs/tagger.js index b203abe5147..2221fb1608a 100644 --- a/packages/dd-trace/src/llmobs/tagger.js +++ b/packages/dd-trace/src/llmobs/tagger.js @@ -93,10 +93,11 @@ class LLMObsTagger { } tagSpanTags (span, tags) { + // new tags will be merged with existing tags try { const currentTags = span.context()._tags[TAGS] if (currentTags) { - Object.assign(tags, currentTags) + Object.assign(tags, JSON.parse(currentTags)) } span.setTag(TAGS, JSON.stringify(tags)) } catch { @@ -110,7 +111,6 @@ class LLMObsTagger { span.setTag(key, data) } else { try { - // this will help showcase unfinished promises being passed in as values span.setTag(key, JSON.stringify(data)) } catch { const type = key === INPUT_VALUE ? 'input' : 'output' @@ -129,33 +129,49 @@ class LLMObsTagger { try { const documents = data.map(document => { if (typeof document === 'string') { - return document + return { text: document } + } + + if (document == null || typeof document !== 'object') { + logger.warn('Documents must be a string, object, or list of objects.') + return undefined } const { text, name, id, score } = document - if (text && typeof text !== 'string') { - logger.warn(`Invalid property found in ${document}: text must be a string.`) + if (typeof text !== 'string') { + logger.warn('Document text must be a string.') return undefined } - if (name && typeof name !== 'string') { - logger.warn(`Invalid property found in ${document}: name must be a string.`) - return undefined + const documentObj = { text } + + if (name) { + if (typeof name !== 'string') { + logger.warn('Document name must be a string.') + return undefined + } + documentObj.name = name } - if (id && typeof id !== 'string') { - logger.warn(`Invalid property found in ${document}: id must be a string.`) - return undefined + if (id) { + if (typeof id !== 'string') { + logger.warn('Document ID must be a string.') + return undefined + } + documentObj.id = id } - if (score && typeof score !== 'number') { - logger.warn(`Invalid property found in ${document}: score must be a number.`) - return undefined + if (score) { + if (typeof score !== 'number') { + logger.warn('Document score must be a number.') + return undefined + } + documentObj.score = score } - return document - }).filter(doc => !!doc) // filter out bad documents? + return documentObj + }).filter(doc => !!doc) span.setTag(key, JSON.stringify(documents)) } catch { @@ -174,28 +190,35 @@ class LLMObsTagger { try { const messages = data.map(message => { if (typeof message === 'string') { - return message + return { content: message } + } + + if (message == null || typeof message !== 'object') { + logger.warn('Messages must be a string, object, or list of objects') + return undefined } const content = message.content || '' const role = message.role if (typeof content !== 'string') { - logger.warn(`Invalid property found in ${message}: content must be a string.`) + logger.warn('Message content must be a string.') return undefined } - message.content = content - - if (role && typeof role !== 'string') { - logger.warn(`Invalid property found in ${message}: role must be a string.`) + if (!role) { + return { content } + } else if (typeof role !== 'string') { + logger.warn('Message role must be a string.') return undefined } - return message - }).filter(msg => !!msg) // filter out bad messages? + return { content, role } + }).filter(msg => !!msg) - span.setTag(key, JSON.stringify(messages)) + if (messages.length) { + span.setTag(key, JSON.stringify(messages)) + } } catch { const type = key === INPUT_MESSAGES ? 'input' : 'output' logger.warn(`Failed to parse ${type} messages.`) diff --git a/packages/dd-trace/test/llmobs/tagger.spec.js b/packages/dd-trace/test/llmobs/tagger.spec.js index d66b5e257f1..db08ddd5544 100644 --- a/packages/dd-trace/test/llmobs/tagger.spec.js +++ b/packages/dd-trace/test/llmobs/tagger.spec.js @@ -184,245 +184,180 @@ describe('tagger', () => { }) }) - describe('tagLLMIO', () => { - beforeEach(() => { - sinon.stub(tagger, '_tagMessages') - }) - - afterEach(() => { - tagger._tagMessages.restore() - }) - - it('tags a span with llm io', () => { - tagger.tagLLMIO(span, { a: 'foo' }, { b: 'bar' }) - expect(tagger._tagMessages).to.have.been.calledTwice - expect(tagger._tagMessages).to.have.been.calledWith(span, { a: 'foo' }, '_ml_obs.meta.input.messages') - expect(tagger._tagMessages).to.have.been.calledWith(span, { b: 'bar' }, '_ml_obs.meta.output.messages') - }) - }) - - describe('tagEmbeddingIO', () => { - beforeEach(() => { - sinon.stub(tagger, '_tagDocuments') - sinon.stub(tagger, '_tagText') - }) - - afterEach(() => { - tagger._tagDocuments.restore() - tagger._tagText.restore() - }) - - it('tags a span with embedding io', () => { - tagger.tagEmbeddingIO(span, { a: 'foo' }, { b: 'bar' }) - expect(tagger._tagDocuments).to.have.been.calledOnce - expect(tagger._tagDocuments).to.have.been.calledWith(span, { a: 'foo' }, '_ml_obs.meta.input.documents') - expect(tagger._tagText).to.have.been.calledOnce - expect(tagger._tagText).to.have.been.calledWith(span, { b: 'bar' }, '_ml_obs.meta.output.value') - }) - }) - - describe('tagRetrievalIO', () => { - beforeEach(() => { - sinon.stub(tagger, '_tagDocuments') - sinon.stub(tagger, '_tagText') - }) - - afterEach(() => { - tagger._tagDocuments.restore() - tagger._tagText.restore() - }) - - it('tags a span with retrieval io', () => { - tagger.tagRetrievalIO(span, { a: 'foo' }, { b: 'bar' }) - expect(tagger._tagDocuments).to.have.been.calledOnce - expect(tagger._tagDocuments).to.have.been.calledWith(span, { b: 'bar' }, '_ml_obs.meta.output.documents') - expect(tagger._tagText).to.have.been.calledOnce - expect(tagger._tagText).to.have.been.calledWith(span, { a: 'foo' }, '_ml_obs.meta.input.value') - }) - }) - - describe('tagTextIO', () => { - beforeEach(() => { - sinon.stub(tagger, '_tagText') - }) - - afterEach(() => { - tagger._tagText.restore() - }) - - it('tags a span with text io', () => { - tagger.tagTextIO(span, { a: 'foo' }, { b: 'bar' }) - expect(tagger._tagText).to.have.been.calledTwice - expect(tagger._tagText).to.have.been.calledWith(span, { a: 'foo' }, '_ml_obs.meta.input.value') - expect(tagger._tagText).to.have.been.calledWith(span, { b: 'bar' }, '_ml_obs.meta.output.value') - }) - }) - - // maybe confirm this one - describe('tagSpanTags', () => {}) - - describe('_tagText', () => { - it('tags a span with text', () => { - tagger._tagText(span, 'my-text', '_ml_obs.meta.input.value') - + describe('tagSpanTags', () => { + it('sets tags on a span', () => { + const tags = { foo: 'bar' } + tagger.tagSpanTags(span, tags) expect(span.context()._tags).to.deep.equal({ - '_ml_obs.meta.input.value': 'my-text' + '_ml_obs.tags': '{"foo":"bar"}' }) }) - it('tags a span with an object', () => { - tagger._tagText(span, { a: 1, b: 2 }, '_ml_obs.meta.input.value') - + it('merges tags so they do not overwrite', () => { + span.context()._tags['_ml_obs.tags'] = '{"a":1}' + const tags = { a: 2, b: 1 } + tagger.tagSpanTags(span, tags) expect(span.context()._tags).to.deep.equal({ - '_ml_obs.meta.input.value': '{"a":1,"b":2}' + '_ml_obs.tags': '{"a":1,"b":1}' }) }) - - it('logs when the value is not JSON serializable', () => { - const data = unserializbleObject() - tagger._tagText(span, data, '_ml_obs.meta.input.value') - expect(logger.warn).to.have.been.calledOnce - }) }) - describe('_tagDocuments', () => { - it('tags a single document object', () => { - const document = { text: 'my-text', name: 'my-name', id: 'my-id', score: 0.5 } - tagger._tagDocuments(span, document, '_ml_obs.meta.input.documents') - - expect(span.context()._tags).to.deep.equal({ - '_ml_obs.meta.input.documents': '[{"text":"my-text","name":"my-name","id":"my-id","score":0.5}]' - }) - }) - - it('tags a single document string', () => { - const document = 'my-text' - tagger._tagDocuments(span, document, '_ml_obs.meta.input.documents') - - expect(span.context()._tags).to.deep.equal({ - '_ml_obs.meta.input.documents': '["my-text"]' - }) - }) + describe('tagLLMIO', () => { + it('tags a span with llm io', () => { + const inputData = [ + 'you are an amazing assistant', + { content: 'hello! my name is foobar' }, + { content: 'I am a robot', role: 'assistant' }, + { content: 'I am a human', role: 'user' } + ] - it('tags a document with a subset of properties', () => { - const document = { text: 'my-text', score: 0.5 } - tagger._tagDocuments(span, document, '_ml_obs.meta.input.documents') + const outputData = 'Nice to meet you, human!' + tagger.tagLLMIO(span, inputData, outputData) expect(span.context()._tags).to.deep.equal({ - '_ml_obs.meta.input.documents': '[{"text":"my-text","score":0.5}]' + '_ml_obs.meta.input.messages': '[{"content":"you are an amazing assistant"},' + + '{"content":"hello! my name is foobar"},{"content":"I am a robot","role":"assistant"},' + + '{"content":"I am a human","role":"user"}]', + '_ml_obs.meta.output.messages': '[{"content":"Nice to meet you, human!"}]' }) }) - it('tags multiple documents', () => { - const documents = [ - { text: 'my-text', name: 'my-name', id: 'my-id', score: 0.5 }, - { text: 'my-text2', name: 'my-name2', id: 'my-id2', score: 0.6 } + it('filters out non-string properties on messages', () => { + const inputData = [ + true, + { content: 5 }, + { content: 'hello', role: 5 }, + 'hi' ] - tagger._tagDocuments(span, documents, '_ml_obs.meta.input.documents') - - expect(span.context()._tags).to.deep.equal({ - '_ml_obs.meta.input.documents': '[{"text":"my-text","name":"my-name"' + - ',"id":"my-id","score":0.5},{"text":"my-text2","name":"my-name2","id":"my-id2","score":0.6}]' - }) - }) - - it('does not include malformed documents', () => { - const documents = [ - { text: 'my-text', name: 'my-name', id: 'my-id', score: 0.5 }, - { text: 'my-text2', score: 'not-a-number' } + const outputData = [ + undefined, + null, + { content: 5 }, + { content: 'goodbye', role: 5 } ] - - tagger._tagDocuments(span, documents, '_ml_obs.meta.input.documents') - + tagger.tagLLMIO(span, inputData, outputData) expect(span.context()._tags).to.deep.equal({ - '_ml_obs.meta.input.documents': '[{"text":"my-text","name":"my-name","id":"my-id","score":0.5}]' + '_ml_obs.meta.input.messages': '[{"content":"hi"}]' }) - }) - it('logs when a document is not JSON serializable', () => { - const documents = [unserializbleObject()] - tagger._tagDocuments(span, documents, '_ml_obs.meta.input.documents') - expect(logger.warn).to.have.been.calledOnce + expect(logger.warn.getCall(0).firstArg).to.equal('Messages must be a string, object, or list of objects') + expect(logger.warn.getCall(1).firstArg).to.equal('Message content must be a string.') + expect(logger.warn.getCall(2).firstArg).to.equal('Message role must be a string.') + expect(logger.warn.getCall(3).firstArg).to.equal('Messages must be a string, object, or list of objects') + expect(logger.warn.getCall(4).firstArg).to.equal('Messages must be a string, object, or list of objects') + expect(logger.warn.getCall(5).firstArg).to.equal('Message content must be a string.') + expect(logger.warn.getCall(6).firstArg).to.equal('Message role must be a string.') }) }) - describe('_tagMessages', () => { - it('tags a single message object', () => { - const message = { content: 'my-content', role: 'my-role' } - tagger._tagMessages(span, message, '_ml_obs.meta.input.messages') - + describe('tagEmbeddingIO', () => { + it('tags a span with embedding io', () => { + const inputData = [ + 'my string document', + { text: 'my object document' }, + { text: 'foo', name: 'bar' }, + { text: 'baz', id: 'qux' }, + { text: 'quux', score: 5 }, + { text: 'foo', name: 'bar', id: 'qux', score: 5 } + ] + const outputData = 'embedded documents' + tagger.tagEmbeddingIO(span, inputData, outputData) expect(span.context()._tags).to.deep.equal({ - '_ml_obs.meta.input.messages': '[{"content":"my-content","role":"my-role"}]' + '_ml_obs.meta.input.documents': '[{"text":"my string document"},{"text":"my object document"},' + + '{"text":"foo","name":"bar"},{"text":"baz","id":"qux"},{"text":"quux","score":5},' + + '{"text":"foo","name":"bar","id":"qux","score":5}]', + '_ml_obs.meta.output.value': 'embedded documents' }) }) - it('tags a single message string', () => { - tagger._tagMessages(span, 'my-message', '_ml_obs.meta.input.messages') - + it('filters out non-string properties on documents', () => { + const inputData = [ + true, + { text: 5 }, + { text: 'foo', name: 5 }, + 'hi', + null, + undefined + ] + const outputData = 'output' + tagger.tagEmbeddingIO(span, inputData, outputData) expect(span.context()._tags).to.deep.equal({ - '_ml_obs.meta.input.messages': '["my-message"]' + '_ml_obs.meta.input.documents': '[{"text":"hi"}]', + '_ml_obs.meta.output.value': 'output' }) - }) - - it('tags a message with only content', () => { - const message = { content: 'my-content' } - tagger._tagMessages(span, message, '_ml_obs.meta.input.messages') - expect(span.context()._tags).to.deep.equal({ - '_ml_obs.meta.input.messages': '[{"content":"my-content"}]' - }) + expect(logger.warn.getCall(0).firstArg).to.equal('Documents must be a string, object, or list of objects.') + expect(logger.warn.getCall(1).firstArg).to.equal('Document text must be a string.') + expect(logger.warn.getCall(2).firstArg).to.equal('Document name must be a string.') + expect(logger.warn.getCall(3).firstArg).to.equal('Documents must be a string, object, or list of objects.') + expect(logger.warn.getCall(4).firstArg).to.equal('Documents must be a string, object, or list of objects.') }) + }) - it('defaults missing content to empty content', () => { - const message = {} - tagger._tagMessages(span, message, '_ml_obs.meta.input.messages') + describe('tagRetrievalIO', () => { + it('tags a span with retrieval io', () => { + const inputData = 'some query' + const outputData = [ + 'result 1', + { text: 'result 2' }, + { text: 'foo', name: 'bar' }, + { text: 'baz', id: 'qux' }, + { text: 'quux', score: 5 }, + { text: 'foo', name: 'bar', id: 'qux', score: 5 } + ] + tagger.tagRetrievalIO(span, inputData, outputData) expect(span.context()._tags).to.deep.equal({ - '_ml_obs.meta.input.messages': '[{"content":""}]' + '_ml_obs.meta.input.value': 'some query', + '_ml_obs.meta.output.documents': '[{"text":"result 1"},{"text":"result 2"},' + + '{"text":"foo","name":"bar"},{"text":"baz","id":"qux"},{"text":"quux","score":5},' + + '{"text":"foo","name":"bar","id":"qux","score":5}]' }) }) - it('filters out messages with non-string content', () => { - const messages = [ - { content: 'my-content' }, - { role: 'my-role', content: 6 } + it('filters out non-string properties on documents', () => { + const inputData = 'some query' + const outputData = [ + true, + { text: 5 }, + { text: 'foo', name: 5 }, + 'hi', + null, + undefined ] - tagger._tagMessages(span, messages, '_ml_obs.meta.input.messages') - + tagger.tagRetrievalIO(span, inputData, outputData) expect(span.context()._tags).to.deep.equal({ - '_ml_obs.meta.input.messages': '[{"content":"my-content"}]' + '_ml_obs.meta.input.value': 'some query', + '_ml_obs.meta.output.documents': '[{"text":"hi"}]' }) - }) - it('filters out messages with non-string role', () => { - const messages = [ - { content: 'my-content', role: 'my-role' }, - { content: 'my-content2', role: 6 } - ] - tagger._tagMessages(span, messages, '_ml_obs.meta.input.messages') + expect(logger.warn.getCall(0).firstArg).to.equal('Documents must be a string, object, or list of objects.') + expect(logger.warn.getCall(1).firstArg).to.equal('Document text must be a string.') + expect(logger.warn.getCall(2).firstArg).to.equal('Document name must be a string.') + expect(logger.warn.getCall(3).firstArg).to.equal('Documents must be a string, object, or list of objects.') + expect(logger.warn.getCall(4).firstArg).to.equal('Documents must be a string, object, or list of objects.') + }) + }) + describe('tagTextIO', () => { + it('tags a span with text io', () => { + const inputData = { some: 'object' } + const outputData = 'some text' + tagger.tagTextIO(span, inputData, outputData) expect(span.context()._tags).to.deep.equal({ - '_ml_obs.meta.input.messages': '[{"content":"my-content","role":"my-role"}]' + '_ml_obs.meta.input.value': '{"some":"object"}', + '_ml_obs.meta.output.value': 'some text' }) }) - it('tags multiple messages', () => { - const messages = [ - { content: 'my-content', role: 'my-role' }, - { content: 'my-content2', role: 'my-role2' } - ] - tagger._tagMessages(span, messages, '_ml_obs.meta.input.messages') - + it('logs when the value is not JSON serializable', () => { + const data = unserializbleObject() + tagger.tagTextIO(span, data, 'output') + expect(logger.warn).to.have.been.calledOnceWith('Failed to parse input value, must be JSON serializable.') expect(span.context()._tags).to.deep.equal({ - '_ml_obs.meta.input.messages': '[{"content":"my-content","role":"my-role"},' + - '{"content":"my-content2","role":"my-role2"}]' + '_ml_obs.meta.output.value': 'output' }) }) - - it('logs when a message is not JSON serializable', () => { - const messages = [unserializbleObject()] - tagger._tagMessages(span, messages, '_ml_obs.meta.input.messages') - expect(logger.warn).to.have.been.calledOnce - }) }) }) From 90e3a168064ad496697e6f74675c8d21b276a6f6 Mon Sep 17 00:00:00 2001 From: Sam Brenner Date: Thu, 26 Sep 2024 16:10:00 -0400 Subject: [PATCH 5/6] enabled, tool call tagging --- packages/dd-trace/src/llmobs/tagger.js | 73 +++++++++++++++++--- packages/dd-trace/test/llmobs/tagger.spec.js | 73 ++++++++++++++++++-- 2 files changed, 131 insertions(+), 15 deletions(-) diff --git a/packages/dd-trace/src/llmobs/tagger.js b/packages/dd-trace/src/llmobs/tagger.js index 2221fb1608a..7c59ef9241e 100644 --- a/packages/dd-trace/src/llmobs/tagger.js +++ b/packages/dd-trace/src/llmobs/tagger.js @@ -36,6 +36,7 @@ class LLMObsTagger { { modelName, modelProvider, sessionId, mlApp, parentLLMObsSpan } = {}, name ) { + if (!this._config.llmobs.enabled) return if (kind) span.setTag(SPAN_TYPE, 'llm') // only mark it as an llm span if it was a valid kind if (name) span.setTag(NAME, name) @@ -198,22 +199,78 @@ class LLMObsTagger { return undefined } - const content = message.content || '' - const role = message.role + const { content = '', role } = message + let toolCalls = message.toolCalls + const messageObj = { content } if (typeof content !== 'string') { logger.warn('Message content must be a string.') return undefined } - if (!role) { - return { content } - } else if (typeof role !== 'string') { - logger.warn('Message role must be a string.') - return undefined + if (role) { + if (typeof role !== 'string') { + logger.warn('Message role must be a string.') + return undefined + } + messageObj.role = role + } + + if (toolCalls) { + if (!Array.isArray(toolCalls)) { + toolCalls = [toolCalls] + } + + const filteredToolCalls = toolCalls.map(toolCall => { + if (typeof toolCall !== 'object') { + logger.warn('Tool call must be an object.') + return undefined + } + + const { name, arguments: args, toolId, type } = toolCall + const toolCallObj = {} + + if (name) { + if (typeof name !== 'string') { + logger.warn('Tool name must be a string.') + return undefined + } + toolCallObj.name = name + } + + if (args) { + if (typeof args !== 'object') { + logger.warn('Tool arguments must be an object.') + return undefined + } + toolCallObj.arguments = args + } + + if (toolId) { + if (typeof toolId !== 'string') { + logger.warn('Tool ID must be a string.') + return undefined + } + toolCallObj.toolId = toolId + } + + if (type) { + if (typeof type !== 'string') { + logger.warn('Tool type must be a string.') + return undefined + } + toolCallObj.type = type + } + + return toolCallObj + }).filter(toolCall => !!toolCall) + + if (filteredToolCalls.length) { + messageObj.tool_calls = filteredToolCalls + } } - return { content, role } + return messageObj }).filter(msg => !!msg) if (messages.length) { diff --git a/packages/dd-trace/test/llmobs/tagger.spec.js b/packages/dd-trace/test/llmobs/tagger.spec.js index db08ddd5544..5d74e11437b 100644 --- a/packages/dd-trace/test/llmobs/tagger.spec.js +++ b/packages/dd-trace/test/llmobs/tagger.spec.js @@ -43,10 +43,17 @@ describe('tagger', () => { './util': util }) - tagger = new Tagger({ llmobs: { mlApp: 'my-default-ml-app' } }) + tagger = new Tagger({ llmobs: { enabled: true, mlApp: 'my-default-ml-app' } }) }) describe('setLLMObsSpanTags', () => { + it('will not set tags if llmobs is not enabled', () => { + tagger = new Tagger({ llmobs: { enabled: false } }) + tagger.setLLMObsSpanTags(span, 'llm') + + expect(span.context()._tags).to.deep.equal({}) + }) + it('tags an llm obs span with basic and default properties', () => { tagger.setLLMObsSpanTags(span, 'workflow') @@ -203,13 +210,14 @@ describe('tagger', () => { }) }) - describe('tagLLMIO', () => { + describe.only('tagLLMIO', () => { it('tags a span with llm io', () => { const inputData = [ 'you are an amazing assistant', { content: 'hello! my name is foobar' }, { content: 'I am a robot', role: 'assistant' }, - { content: 'I am a human', role: 'user' } + { content: 'I am a human', role: 'user' }, + {} ] const outputData = 'Nice to meet you, human!' @@ -218,12 +226,12 @@ describe('tagger', () => { expect(span.context()._tags).to.deep.equal({ '_ml_obs.meta.input.messages': '[{"content":"you are an amazing assistant"},' + '{"content":"hello! my name is foobar"},{"content":"I am a robot","role":"assistant"},' + - '{"content":"I am a human","role":"user"}]', + '{"content":"I am a human","role":"user"},{"content":""}]', '_ml_obs.meta.output.messages': '[{"content":"Nice to meet you, human!"}]' }) }) - it('filters out non-string properties on messages', () => { + it('filters out malformed properties on messages', () => { const inputData = [ true, { content: 5 }, @@ -249,6 +257,57 @@ describe('tagger', () => { expect(logger.warn.getCall(5).firstArg).to.equal('Message content must be a string.') expect(logger.warn.getCall(6).firstArg).to.equal('Message role must be a string.') }) + + describe('tagging tool calls appropriately', () => { + it('tags a span with tool calls', () => { + const inputData = [ + { content: 'hello', toolCalls: [{ name: 'tool1' }, { name: 'tool2', arguments: { a: 1, b: 2 } }] }, + { content: 'goodbye', toolCalls: [{ name: 'tool3' }] } + ] + const outputData = [ + { content: 'hi', toolCalls: [{ name: 'tool4' }] } + ] + + tagger.tagLLMIO(span, inputData, outputData) + expect(span.context()._tags).to.deep.equal({ + '_ml_obs.meta.input.messages': '[{"content":"hello","tool_calls":[{"name":"tool1"},' + + '{"name":"tool2","arguments":{"a":1,"b":2}}]},' + + '{"content":"goodbye","tool_calls":[{"name":"tool3"}]}]', + '_ml_obs.meta.output.messages': '[{"content":"hi","tool_calls":[{"name":"tool4"}]}]' + }) + }) + + it('filters out malformed tool calls', () => { + const inputData = [ + { content: 'a', toolCalls: 5 }, // tool calls must be objects + { content: 'b', toolCalls: [5] }, // tool calls must be objects + { content: 'c', toolCalls: [{ name: 5 }] }, // tool name must be a string + { content: 'd', toolCalls: [{ arguments: 5 }] }, // tool arguments must be an object + { content: 'e', toolCalls: [{ toolId: 5 }] }, // tool id must be a string + { content: 'f', toolCalls: [{ type: 5 }] }, // tool type must be a string + { + content: 'g', + toolCalls: [ + { name: 'tool1', arguments: 5 }, { name: 'tool2' } // second tool call should be tagged + ] + } // tool arguments must be an object + ] + + tagger.tagLLMIO(span, inputData, undefined) + expect(span.context()._tags).to.deep.equal({ + '_ml_obs.meta.input.messages': '[{"content":"a"},{"content":"b"},{"content":"c"},' + + '{"content":"d"},{"content":"e"},{"content":"f"},{"content":"g","tool_calls":[{"name":"tool2"}]}]' + }) + + expect(logger.warn.getCall(0).firstArg).to.equal('Tool call must be an object.') + expect(logger.warn.getCall(1).firstArg).to.equal('Tool call must be an object.') + expect(logger.warn.getCall(2).firstArg).to.equal('Tool name must be a string.') + expect(logger.warn.getCall(3).firstArg).to.equal('Tool arguments must be an object.') + expect(logger.warn.getCall(4).firstArg).to.equal('Tool ID must be a string.') + expect(logger.warn.getCall(5).firstArg).to.equal('Tool type must be a string.') + expect(logger.warn.getCall(6).firstArg).to.equal('Tool arguments must be an object.') + }) + }) }) describe('tagEmbeddingIO', () => { @@ -271,7 +330,7 @@ describe('tagger', () => { }) }) - it('filters out non-string properties on documents', () => { + it('filters out malformed properties on documents', () => { const inputData = [ true, { text: 5 }, @@ -316,7 +375,7 @@ describe('tagger', () => { }) }) - it('filters out non-string properties on documents', () => { + it('filters out malformed properties on documents', () => { const inputData = 'some query' const outputData = [ true, From 3083abf4ef625db3070bebd8bd0dcb7cfe79c740 Mon Sep 17 00:00:00 2001 From: Sam Brenner Date: Thu, 26 Sep 2024 17:13:41 -0400 Subject: [PATCH 6/6] remove .only --- packages/dd-trace/test/llmobs/tagger.spec.js | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/dd-trace/test/llmobs/tagger.spec.js b/packages/dd-trace/test/llmobs/tagger.spec.js index 5d74e11437b..935389cd985 100644 --- a/packages/dd-trace/test/llmobs/tagger.spec.js +++ b/packages/dd-trace/test/llmobs/tagger.spec.js @@ -210,7 +210,7 @@ describe('tagger', () => { }) }) - describe.only('tagLLMIO', () => { + describe('tagLLMIO', () => { it('tags a span with llm io', () => { const inputData = [ 'you are an amazing assistant',