add decorate to noop

DataDog · Oct 11, 2024 · a4001d5 · a4001d5
1 parent db8e2eb
commit a4001d5
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 22 deletions.
diff --git a/packages/dd-trace/src/llmobs/noop.js b/packages/dd-trace/src/llmobs/noop.js
@@ -19,40 +19,57 @@ class NoopLLMObs {
 
   submitEvaluation (llmobsSpanContext, options) {}
 
-  startSpan (kind, options) {
-    return this._tracer.startSpan(kind, options)
+  startSpan (options = {}) {
+    const name = options.name || options.kind
+    return this._tracer.startSpan(name, options)
   }
 
-  trace (kind, options, fn) {
-    if (!fn) {
-      fn = options
-      options = {}
-    }
-
+  trace (fn, options = {}) {
     if (typeof fn !== 'function') return
 
-    options = options || {}
+    const name = options.name || options.kind || fn.name
 
-    return this._tracer.trace(kind, options, fn)
+    return this._tracer.trace(name, options, fn)
   }
 
-  wrap (kind, options, fn) {
-    if (!fn) {
-      fn = options
-      options = {}
-    }
-
+  wrap (fn, options = {}) {
     if (typeof fn !== 'function') return fn
 
-    options = options || {}
+    const name = options.name || options.kind || fn.name
 
-    return this._tracer.wrap(kind, options, fn)
+    return this._tracer.wrap(name, options, fn)
   }
 
-  decorate (kind, options) {
+  decorate (options = {}) {
     const llmobs = this
-    return function (target) {
-      return llmobs.wrap(kind, options, target)
+    return function (target, ctxOrPropertyKey, descriptor) {
+      if (!ctxOrPropertyKey) return target
+      if (typeof ctxOrPropertyKey === 'object') {
+        const ctx = ctxOrPropertyKey
+        if (ctx.kind !== 'method') return target
+
+        return llmobs.wrap(target, { name: ctx.name, ...options })
+      } else {
+        const propertyKey = ctxOrPropertyKey
+        if (descriptor) {
+          if (typeof descriptor.value !== 'function') return descriptor
+
+          const original = descriptor.value
+          descriptor.value = llmobs.wrap(original, { name: propertyKey, ...options })
+
+          return descriptor
+        } else {
+          if (typeof target[propertyKey] !== 'function') return target[propertyKey]
+
+          const original = target[propertyKey]
+          Object.defineProperty(target, propertyKey, {
+            ...Object.getOwnPropertyDescriptor(target, propertyKey),
+            value: llmobs.wrap(original, { name: propertyKey, ...options })
+          })
+
+          return target
+        }
+      }
     }
   }
 

diff --git a/packages/dd-trace/src/llmobs/sdk.js b/packages/dd-trace/src/llmobs/sdk.js
@@ -484,7 +484,7 @@ class LLMObs {
 
           return descriptor
         } else {
-          if (typeof target[propertyKey] !== 'function') return target
+          if (typeof target[propertyKey] !== 'function') return target[propertyKey]
 
           const original = target[propertyKey]
           Object.defineProperty(target, propertyKey, {