@senzops/apm-node 1.2.8 → 1.3.1

package/src/instrumentation/cohere.ts

@@ -0,0 +1,227 @@
+import { SenzorOptions } from '../core/types';
+import { hookRequire } from './hook';
+import { patchMethod } from './patch';
+import { runWithCapturedSpan, startCapturedSpan } from './span';
+
+// ---------------------------------------------------------------------------
+// Cohere SDK Instrumentation
+//
+// Instruments the `cohere-ai` npm package for Cohere's NLP models
+// (Command, Embed, Rerank, Classify, Summarize, etc.).
+//
+// Patches CohereClient.prototype (or CohereClientV2.prototype) methods:
+//   - chat()          — conversational generation (Command R/R+)
+//   - chatStream()    — streaming chat
+//   - generate()      — text generation (legacy)
+//   - embed()         — embedding generation
+//   - rerank()        — semantic reranking
+//   - classify()      — text classification
+//   - summarize()     — text summarization
+//   - tokenize()      — tokenization
+//   - detokenize()    — detokenization
+//
+// Captured attributes (OTel GenAI semantic conventions):
+//   - gen_ai.system: 'cohere'
+//   - gen_ai.request.model: command-r-plus, embed-english-v3.0, etc.
+//   - gen_ai.operation.name: chat, generate, embed, rerank, etc.
+//   - gen_ai.usage.input_tokens: billed input tokens
+//   - gen_ai.usage.output_tokens: billed output tokens
+//   - gen_ai.response.finish_reason: COMPLETE, MAX_TOKENS, etc.
+// ---------------------------------------------------------------------------
+
+/** Methods to instrument with metadata extractors. */
+const METHODS: {
+  name: string;
+  operation: string;
+  getModel: (args: any[]) => string | undefined;
+  extractUsage: (result: any) => Record<string, any>;
+}[] = [
+  {
+    name: 'chat',
+    operation: 'chat',
+    getModel: (args) => args[0]?.model,
+    extractUsage: (result) => {
+      const meta: Record<string, any> = {};
+      // Cohere v2 chat response
+      if (result?.meta?.billedUnits) {
+        meta['gen_ai.usage.input_tokens'] = result.meta.billedUnits.inputTokens;
+        meta['gen_ai.usage.output_tokens'] = result.meta.billedUnits.outputTokens;
+      }
+      // Cohere v1 chat response
+      if (result?.meta?.tokens) {
+        meta['gen_ai.usage.input_tokens'] = meta['gen_ai.usage.input_tokens'] || result.meta.tokens.inputTokens;
+        meta['gen_ai.usage.output_tokens'] = meta['gen_ai.usage.output_tokens'] || result.meta.tokens.outputTokens;
+      }
+      if (result?.finishReason || result?.finish_reason) {
+        meta['gen_ai.response.finish_reason'] = result.finishReason || result.finish_reason;
+      }
+      return meta;
+    },
+  },
+  {
+    name: 'chatStream',
+    operation: 'chat.stream',
+    getModel: (args) => args[0]?.model,
+    extractUsage: () => ({}), // Stream — usage comes in final event
+  },
+  {
+    name: 'generate',
+    operation: 'generate',
+    getModel: (args) => args[0]?.model,
+    extractUsage: (result) => {
+      const meta: Record<string, any> = {};
+      if (result?.meta?.billedUnits) {
+        meta['gen_ai.usage.input_tokens'] = result.meta.billedUnits.inputTokens;
+        meta['gen_ai.usage.output_tokens'] = result.meta.billedUnits.outputTokens;
+      }
+      return meta;
+    },
+  },
+  {
+    name: 'embed',
+    operation: 'embed',
+    getModel: (args) => args[0]?.model,
+    extractUsage: (result) => {
+      const meta: Record<string, any> = {};
+      if (result?.meta?.billedUnits) {
+        meta['gen_ai.usage.input_tokens'] = result.meta.billedUnits.inputTokens;
+      }
+      return meta;
+    },
+  },
+  {
+    name: 'rerank',
+    operation: 'rerank',
+    getModel: (args) => args[0]?.model,
+    extractUsage: (result) => {
+      const meta: Record<string, any> = {};
+      if (result?.meta?.billedUnits) {
+        meta['gen_ai.usage.input_tokens'] = result.meta.billedUnits.searchUnits;
+      }
+      meta['cohere.results_count'] = result?.results?.length;
+      return meta;
+    },
+  },
+  {
+    name: 'classify',
+    operation: 'classify',
+    getModel: (args) => args[0]?.model,
+    extractUsage: () => ({}),
+  },
+  {
+    name: 'summarize',
+    operation: 'summarize',
+    getModel: (args) => args[0]?.model,
+    extractUsage: (result) => {
+      const meta: Record<string, any> = {};
+      if (result?.meta?.billedUnits) {
+        meta['gen_ai.usage.input_tokens'] = result.meta.billedUnits.inputTokens;
+        meta['gen_ai.usage.output_tokens'] = result.meta.billedUnits.outputTokens;
+      }
+      return meta;
+    },
+  },
+  {
+    name: 'tokenize',
+    operation: 'tokenize',
+    getModel: (args) => args[0]?.model,
+    extractUsage: (result) => ({
+      'cohere.token_count': result?.tokens?.length,
+    }),
+  },
+  {
+    name: 'detokenize',
+    operation: 'detokenize',
+    getModel: (args) => args[0]?.model,
+    extractUsage: () => ({}),
+  },
+];
+
+// ---------------------------------------------------------------------------
+// CohereClient patching
+// ---------------------------------------------------------------------------
+
+const patchCohereClient = (proto: any, clientName: string, options?: SenzorOptions) => {
+  if (!proto) return;
+
+  for (const methodConfig of METHODS) {
+    if (typeof proto[methodConfig.name] !== 'function') continue;
+
+    patchMethod(
+      proto,
+      methodConfig.name,
+      `senzor.cohere.${clientName}.${methodConfig.name}`,
+      (original) =>
+        function patchedCohereMethod(this: any, ...args: any[]) {
+          const model = methodConfig.getModel(args);
+
+          const spanName = model
+            ? `Cohere ${methodConfig.operation} ${model}`
+            : `Cohere ${methodConfig.operation}`;
+
+          const span = startCapturedSpan(
+            spanName,
+            'http',
+            {
+              'gen_ai.system': 'cohere',
+              'gen_ai.operation.name': methodConfig.operation,
+              'gen_ai.request.model': model,
+              library: 'cohere',
+            },
+            options
+          );
+
+          if (!span) return original.apply(this, args);
+
+          return runWithCapturedSpan(span, () => {
+            try {
+              const result = original.apply(this, args);
+
+              if (result && typeof result.then === 'function') {
+                return result.then(
+                  (value: any) => {
+                    span.end(0, methodConfig.extractUsage(value));
+                    return value;
+                  },
+                  (error: any) => {
+                    span.end(error?.statusCode || error?.status || 500, {
+                      'error.message': error?.message,
+                      'error.type': error?.name || 'CohereError',
+                    });
+                    throw error;
+                  }
+                );
+              }
+
+              span.end(0);
+              return result;
+            } catch (error: any) {
+              span.end(500, { 'error.message': error?.message });
+              throw error;
+            }
+          });
+        }
+    );
+  }
+};
+
+// ---------------------------------------------------------------------------
+// Public API
+// ---------------------------------------------------------------------------
+
+export const instrumentCohere = (options?: SenzorOptions) => {
+  hookRequire('cohere-ai', (exports: any) => {
+    // CohereClient (v1)
+    if (exports?.CohereClient?.prototype) {
+      patchCohereClient(exports.CohereClient.prototype, 'client', options);
+    }
+    // CohereClientV2 (v2)
+    if (exports?.CohereClientV2?.prototype) {
+      patchCohereClient(exports.CohereClientV2.prototype, 'clientV2', options);
+    }
+    // Default export
+    if (exports?.default?.prototype) {
+      patchCohereClient(exports.default.prototype, 'default', options);
+    }
+  });
+};

package/src/instrumentation/connect.ts

@@ -0,0 +1,200 @@
+import { SenzorOptions } from '../core/types';
+import { hookRequire } from './hook';
+import { patchMethod } from './patch';
+import { runWithCapturedSpan, startCapturedSpan } from './span';
+
+// ---------------------------------------------------------------------------
+// Connect Instrumentation
+//
+// Instruments the `connect` middleware framework — the foundation that
+// Express was originally built on. Many production apps still use
+// Connect directly for lightweight HTTP services.
+//
+// Patches the connect app's use() method to wrap every middleware
+// function with a span capturing execution time and errors.
+//
+// Connect middleware signature: (req, res, next) or (err, req, res, next)
+//
+// Captured attributes:
+//   - connect.type: 'middleware'
+//   - connect.name: middleware function name
+//   - connect.route: mount path (if provided)
+//   - framework: 'connect'
+// ---------------------------------------------------------------------------
+
+// ---------------------------------------------------------------------------
+// Middleware wrapping
+// ---------------------------------------------------------------------------
+
+const wrapMiddleware = (
+  fn: Function,
+  route: string,
+  options?: SenzorOptions
+): Function => {
+  if (typeof fn !== 'function') return fn;
+  if ((fn as any).__senzorWrapped) return fn;
+
+  const middlewareName = fn.name || 'anonymous';
+  const isErrorHandler = fn.length >= 4;
+
+  let wrapped: Function;
+
+  if (isErrorHandler) {
+    // Error-handling middleware: (err, req, res, next)
+    wrapped = function wrappedConnectErrorMiddleware(
+      this: any, err: any, req: any, res: any, next: any
+    ) {
+      const span = startCapturedSpan(
+        `Connect error ${middlewareName}`,
+        'function',
+        {
+          'connect.type': 'error_middleware',
+          'connect.name': middlewareName,
+          'connect.route': route,
+          framework: 'connect',
+        },
+        options
+      );
+
+      if (!span) return fn.call(this, err, req, res, next);
+
+      return runWithCapturedSpan(span, () => {
+        const wrappedNext = function (...args: any[]) {
+          const hasError = args.length > 0 && args[0] instanceof Error;
+          span.end(hasError ? 500 : 0, hasError ? { 'error.message': args[0].message } : {});
+          return next?.(...args);
+        };
+
+        try {
+          const result = fn.call(this, err, req, res, wrappedNext);
+          if (result && typeof result.then === 'function') {
+            return result.catch((error: any) => {
+              span.end(500, { 'error.message': error?.message });
+              throw error;
+            });
+          }
+          return result;
+        } catch (error: any) {
+          span.end(500, { 'error.message': error?.message });
+          throw error;
+        }
+      });
+    };
+  } else {
+    // Standard middleware: (req, res, next)
+    wrapped = function wrappedConnectMiddleware(
+      this: any, req: any, res: any, next: any
+    ) {
+      const span = startCapturedSpan(
+        `Connect ${middlewareName}`,
+        'function',
+        {
+          'connect.type': 'middleware',
+          'connect.name': middlewareName,
+          'connect.route': route,
+          'http.route': route !== '/' ? route : undefined,
+          framework: 'connect',
+        },
+        options
+      );
+
+      if (!span) return fn.call(this, req, res, next);
+
+      return runWithCapturedSpan(span, () => {
+        const wrappedNext = function (...args: any[]) {
+          const hasError = args.length > 0 && args[0] instanceof Error;
+          span.end(hasError ? 500 : 0, hasError ? { 'error.message': args[0].message } : {});
+          return next?.(...args);
+        };
+
+        try {
+          const result = fn.call(this, req, res, wrappedNext);
+          if (result && typeof result.then === 'function') {
+            return result.catch((error: any) => {
+              span.end(500, { 'error.message': error?.message });
+              throw error;
+            });
+          }
+          return result;
+        } catch (error: any) {
+          span.end(500, { 'error.message': error?.message });
+          throw error;
+        }
+      });
+    };
+  }
+
+  // Preserve original function length for Connect's error handler detection
+  Object.defineProperty(wrapped, 'length', { value: fn.length });
+  (wrapped as any).__senzorWrapped = true;
+
+  return wrapped;
+};
+
+// ---------------------------------------------------------------------------
+// app.use() patching
+// ---------------------------------------------------------------------------
+
+const patchConnectApp = (connectModule: any, options?: SenzorOptions) => {
+  if (typeof connectModule !== 'function') return;
+
+  // connect() returns an app. We need to patch the app's prototype.
+  // Connect apps have use() on their prototype chain via proto.
+
+  // Create a probe app to get the prototype
+  let appProto: any;
+
+  try {
+    const app = connectModule();
+    appProto = Object.getPrototypeOf(app);
+  } catch { }
+
+  if (!appProto) return;
+
+  // Patch use()
+  if (typeof appProto.use === 'function') {
+    patchMethod(
+      appProto,
+      'use',
+      'senzor.connect.app.use',
+      (original) =>
+        function patchedUse(this: any, ...args: any[]) {
+          // use() accepts:
+          //   use(fn)
+          //   use(route, fn)
+          //   use(route, fn1, fn2, ...)
+
+          let route = '/';
+          let startIdx = 0;
+
+          if (typeof args[0] === 'string') {
+            route = args[0];
+            startIdx = 1;
+          }
+
+          // Wrap all function arguments
+          for (let i = startIdx; i < args.length; i++) {
+            if (typeof args[i] === 'function') {
+              args[i] = wrapMiddleware(args[i], route, options);
+            }
+          }
+
+          return original.apply(this, args);
+        }
+    );
+  }
+};
+
+// ---------------------------------------------------------------------------
+// Public API
+// ---------------------------------------------------------------------------
+
+export const instrumentConnect = (options?: SenzorOptions) => {
+  hookRequire('connect', (exports: any) => {
+    patchConnectApp(exports, options);
+
+    if (exports?.default) {
+      patchConnectApp(exports.default, options);
+    }
+  });
+};