feat: add OITracer to autoinstrumentations providing support for attr…

…ibute masking and context attribute inheritance (#1078)

Co-authored-by: Tony Powell <[email protected]>

js/.changeset/khaki-students-attend.md

@@ -0,0 +1,6 @@
+---
+"@arizeai/openinference-instrumentation-langchain": major
+"@arizeai/openinference-instrumentation-openai": major
+---
+
+add support for trace config to OpenAI and LangChain auto instrumentors to allow for attribute masking on spans

js/DEVELOPMENT.md

@@ -1,5 +1,16 @@
 ## JavaScript Development
 
+- [Setup](#setup)
+- [Testing](#testing)
+- [Creating an Instrumentor](#creating-an-instrumentor)
+  - [Minimum Feature Set](#minimum-feature-set)
+    - [Suppress Tracing](#suppress-tracing)
+    - [Context Attribute Propagation](#context-attribute-propagation)
+    - [Trace Configuration](#trace-configuration)
+    - [Testing](#testing-1)
+- [Changesets](#changesets)
+- [Publishing](#publishing)
+
 The development guide for the JavaScript packages in this repo.
 
 This project and its packages are built using the following tools:
@@ -51,6 +62,66 @@ pnpm run -r test
 > [!NOTE]
 > The tests in this repo use `jest` but it's auto-mocking feature can cause issues since instrumentation relies on it running first before the package is imported in user-code. For the tests you may have to manually set the instrumented module manually (e.x.`instrumentation._modules[0].moduleExports = module`)
 
+## Creating an Instrumentor
+
+To keep our instrumentors up to date and in sync there is a set of features that each instrumentor must contain. Most of these features are implemented in our [openinference-core package](./packages/openinference-core/) and will be handled by the [OITracer](./packages/openinference-core/src/trace/trace-config/OITracer.ts) and underlying [OISpan](./packages/openinference-core/src/trace/trace-config/OISpan.ts) so it's important to use the OITracer in your instrumentations.
+
+To use the OITracer within your instrumentors, make sure to set the OITracer as a private property on your instrumentor class. Then be sure to use it anytime you need to create a span (or do anything else with the `tracer`). Example:
+
+```typescript
+export class MyInstrumentation extends InstrumentationBase<
+  typeof moduleToInstrument
+> {
+  private oiTracer: OITracer;
+  constructor({
+    instrumentationConfig,
+    traceConfig,
+  }: {
+    instrumentationConfig?: InstrumentationConfig;
+    traceConfig?: TraceConfigOptions;
+  } = {}) {
+    super(
+      "@arizeai/openinference-instrumentation-example-module",
+      VERSION,
+      Object.assign({}, instrumentationConfig),
+    );
+    this.oiTracer = new OITracer({ tracer: this.tracer, traceConfig });
+  }
+  // ...
+  // Use this.oiTracer not this.tracer in your code to create spans
+  const span = this.oiTracer.startSpan("my-span");
+  span.end()
+}
+```
+
+### Minimum Feature Set
+
+Each instrumentation must contain the following features:
+
+#### Suppress Tracing
+
+Every instrumentor must allow tracing to be suppressed or disabled.
+
+In JS/TS tracing suppression is controlled by a context attribute see [suppress-tracing.ts](https://github.com/open-telemetry/opentelemetry-js/blob/55a1fc88d84b22c08e6a19eff71875e15377b781/packages/opentelemetry-core/src/trace/suppress-tracing.ts#L23) from opentelemetry-js. This context key must be respected in each instrumentation. To check for this key and block tracing see our [openai-instrumentation](./packages/openinference-instrumentation-openai/src/instrumentation.ts#69).
+
+Every instrumentation must also be able to be disabled. The `disable` method is inherited from the `InstrumentationBase` class and does not have to be implemented. To ensure that your instrumentation can be properly disabled you just need to properly implement the `unpatch` method on your instrumentation.
+
+#### Context Attribute Propagation
+
+There are a number of situations in which a user may want to attach a specific attribute to every span that gets created within a particular block or scope. For example a user may want to ensure that every span created has a user or session ID attached to it. We achieve this by allowing users to set attributes on [context](https://opentelemetry.io/docs/specs/otel/context/). Our instrumentors must respect these attributes and correctly propagate them to each span.
+
+This fetching and propagation is controlled by our [OITracer](./packages/openinference-core/src/trace/trace-config/OITracer.ts#117) and [context attributes](./packages/openinference-core/src/trace/contextAttributes.ts) from our core package. See the example above to properly use the OITracer in your instrumentor to ensure context attributes are repected.
+
+#### Trace Configuration
+
+In some situations, users may want to control what data gets added to a span. We allow them to do this via [trace config](./packages/openinference-core/src/trace/trace-config/). Our trace config allows users to mask certain fields on a span to prevent sensitive information from leaving their system.
+
+As with context attribute propagation, this is controlled by our OITracer and [OISpan](./packages/openinference-core/src/trace/trace-config/OISpan.ts#21). See the example above to properly use the OITracer in your instrumentor to ensure the trace config is respected.
+
+#### Testing
+
+In addition to any additional testing you do for your instrumentor, it's important to write tests specifically for the features above. This ensures that all of our instrumentors have the same core set of functionality and can help to catch up-stream bugs in our core package.
+
 ## Changesets
 
 The changes to the packages managed by this repo are tracked via [changesets](https://pnpm.io/using-changesets). Changesets are similar to semantic commits in that they describe the changes made to the codebase. However, changesets track changes to all the packages by committing `changesets` to the `.changeset` directory. If you make a change to a package, you should create a changeset for it via:
@@ -59,9 +130,13 @@ The changes to the packages managed by this repo are tracked via [changesets](ht
 pnpm changeset
 ```
 
+and commit it in your pr.
+
 A changeset is an intent to release a set of packages at particular [semver bump types](https://semver.org/) with a summary of the changes made.
 
-For a detailed explanation of changesets, consult [this documentation])(https://github.com/changesets/changesets/blob/main/docs/detailed-explanation.md)
+Once your pr is merged, Github Actions will create a release PR like [this](https://github.com/Arize-ai/openinference/pull/994). Once the release pr is merged, new versions of any changed packages will be published to npm.
+
+For a detailed explanation of changesets, consult [this documentation](https://github.com/changesets/changesets/blob/main/docs/detailed-explanation.md)
 
 ## Publishing
 

js/packages/openinference-core/test/trace-config/OITracer.test.ts

@@ -10,11 +10,21 @@ import {
   Tracer,
 } from "@opentelemetry/api";
 import { AsyncHooksContextManager } from "@opentelemetry/context-async-hooks";
+import { NodeTracerProvider } from "@opentelemetry/sdk-trace-node";
+import {
+  InMemorySpanExporter,
+  SimpleSpanProcessor,
+} from "@opentelemetry/sdk-trace-base";
+
+const tracerProvider = new NodeTracerProvider();
+tracerProvider.register();
 
 describe("OITracer", () => {
   let mockTracer: jest.Mocked<Tracer>;
   let mockSpan: jest.Mocked<Span>;
   let contextManager: ContextManager;
+  const memoryExporter = new InMemorySpanExporter();
+  tracerProvider.addSpanProcessor(new SimpleSpanProcessor(memoryExporter));
 
   beforeEach(() => {
     contextManager = new AsyncHooksContextManager().enable();
@@ -31,7 +41,9 @@ describe("OITracer", () => {
       }),
     } as unknown as jest.Mocked<Tracer>;
   });
-
+  beforeEach(() => {
+    memoryExporter.reset();
+  });
   afterEach(() => {
     context.disable();
   });
@@ -95,7 +107,6 @@ describe("OITracer", () => {
         },
       );
     });
-    it("should correctly nest spans", () => {});
   });
 
   describe("startActiveSpan", () => {
@@ -206,5 +217,42 @@ describe("OITracer", () => {
         },
       );
     });
+    it("should properly nest spans", () => {
+      const tracer = tracerProvider.getTracer("test");
+
+      const oiTracer = new OITracer({
+        tracer,
+        traceConfig: {
+          hideInputs: true,
+        },
+      });
+
+      oiTracer.startActiveSpan("parent", (parentSpan) => {
+        const childSpan = oiTracer.startSpan("child");
+        childSpan.end();
+        parentSpan.end();
+      });
+
+      oiTracer.startSpan("parent2").end();
+
+      const spans = memoryExporter.getFinishedSpans();
+      const parentSpan = spans.find((span) => span.name === "parent");
+      const childSpan = spans.find((span) => span.name === "child");
+      const parent2 = spans.find((span) => span.name === "parent2");
+
+      expect(parentSpan).toBeDefined();
+      expect(childSpan).toBeDefined();
+      const parentSpanId = parentSpan?.spanContext().spanId;
+      expect(parentSpanId).toBeDefined();
+      const childSpanParentId = childSpan?.parentSpanId;
+      expect(childSpanParentId).toBeDefined();
+      expect(childSpanParentId).toBe(parentSpanId);
+      expect(childSpan?.spanContext().traceId).toBe(
+        parentSpan?.spanContext().traceId,
+      );
+      expect(parent2?.spanContext().traceId).not.toBe(
+        childSpan?.spanContext().traceId,
+      );
+    });
   });
 });

js/packages/openinference-instrumentation-langchain/src/instrumentation.ts

@@ -10,6 +10,7 @@ import {
 import { VERSION } from "./version";
 import { diag } from "@opentelemetry/api";
 import { addTracerToHandlers } from "./instrumentationUtils";
+import { OITracer, TraceConfigOptions } from "@arizeai/openinference-core";
 
 const MODULE_NAME = "@langchain/core/callbacks";
 
@@ -30,13 +31,38 @@ type CallbackManagerModule =
   | typeof CallbackManagerModuleV01
   | typeof CallbackManagerModuleV02;
 
+/**
+ * An auto instrumentation class for LangChain that creates {@link https://github.com/Arize-ai/openinference/blob/main/spec/semantic_conventions.md|OpenInference} Compliant spans for LangChain
+ * @param instrumentationConfig The config for the instrumentation @see {@link InstrumentationConfig}
+ * @param traceConfig The OpenInference trace configuration. Can be used to mask or redact sensitive information on spans. @see {@link TraceConfigOptions}
+ */
 export class LangChainInstrumentation extends InstrumentationBase<CallbackManagerModule> {
-  constructor(config?: InstrumentationConfig) {
+  private oiTracer: OITracer;
+
+  constructor({
+    instrumentationConfig,
+    traceConfig,
+  }: {
+    /**
+     * The config for the instrumentation
+     * @see {@link InstrumentationConfig}
+     */
+    instrumentationConfig?: InstrumentationConfig;
+    /**
+     * The OpenInference trace configuration. Can be used to mask or redact sensitive information on spans.
+     * @see {@link TraceConfigOptions}
+     */
+    traceConfig?: TraceConfigOptions;
+  } = {}) {
     super(
       "@arizeai/openinference-instrumentation-langchain",
       VERSION,
-      Object.assign({}, config),
+      Object.assign({}, instrumentationConfig),
     );
+    this.oiTracer = new OITracer({
+      tracer: this.tracer,
+      traceConfig,
+    });
   }
 
   manuallyInstrument(module: CallbackManagerModule) {
@@ -90,7 +116,7 @@ export class LangChainInstrumentation extends InstrumentationBase<CallbackManage
         ) {
           const inheritableHandlers = args[0];
           const newInheritableHandlers = addTracerToHandlers(
-            instrumentation.tracer,
+            instrumentation.oiTracer,
             inheritableHandlers,
           );
           args[0] = newInheritableHandlers;
@@ -108,7 +134,7 @@ export class LangChainInstrumentation extends InstrumentationBase<CallbackManage
         ) {
           const handlers = args[0];
           const newHandlers = addTracerToHandlers(
-            instrumentation.tracer,
+            instrumentation.oiTracer,
             handlers,
           );
           args[0] = newHandlers;
@@ -122,7 +148,7 @@ export class LangChainInstrumentation extends InstrumentationBase<CallbackManage
       // This can fail if the module is made immutable via the runtime or bundler
       module.openInferencePatched = true;
     } catch (e) {
-      diag.warn(`Failed to set ${MODULE_NAME} patched flag on the module`, e);
+      diag.debug(`Failed to set ${MODULE_NAME} patched flag on the module`, e);
     }
 
     return module;

js/packages/openinference-instrumentation-langchain/src/instrumentationUtils.ts

@@ -1,7 +1,7 @@
 import type * as CallbackManagerModuleV02 from "@langchain/core/callbacks/manager";
 import type * as CallbackManagerModuleV01 from "@langchain/coreV0.1/callbacks/manager";
-import { Tracer } from "@opentelemetry/api";
 import { LangChainTracer } from "./tracer";
+import { OITracer } from "@arizeai/openinference-core";
 
 /**
  * Adds the {@link LangChainTracer} to the callback handlers if it is not already present
@@ -16,15 +16,15 @@ import { LangChainTracer } from "./tracer";
  * We support both versions and our tracer is compatible with either as it will extend the BaseTracer from the installed version which will be the same as the version of handlers passed in here
  */
 export function addTracerToHandlers(
-  tracer: Tracer,
+  tracer: OITracer,
   handlers?: CallbackManagerModuleV01.Callbacks,
 ): CallbackManagerModuleV01.Callbacks;
 export function addTracerToHandlers(
-  tracer: Tracer,
+  tracer: OITracer,
   handlers?: CallbackManagerModuleV02.Callbacks,
 ): CallbackManagerModuleV02.Callbacks;
 export function addTracerToHandlers(
-  tracer: Tracer,
+  tracer: OITracer,
   handlers?:
     | CallbackManagerModuleV01.Callbacks
     | CallbackManagerModuleV02.Callbacks,

js/packages/openinference-instrumentation-langchain/src/tracer.ts

@@ -1,6 +1,5 @@
 import { BaseTracer, Run } from "@langchain/core/tracers/base";
 import {
-  Tracer,
   SpanKind,
   Span,
   context,
@@ -23,16 +22,17 @@ import {
   safelyFormatToolCalls,
   safelyGetOpenInferenceSpanKindFromRunType,
 } from "./utils";
+import { OITracer } from "@arizeai/openinference-core";
 
 type RunWithSpan = {
   run: Run;
   span: Span;
 };
 
 export class LangChainTracer extends BaseTracer {
-  private tracer: Tracer;
+  private tracer: OITracer;
   private runs: Record<string, RunWithSpan | undefined> = {};
-  constructor(tracer: Tracer) {
+  constructor(tracer: OITracer) {
     super();
     this.tracer = tracer;
   }

js/packages/openinference-instrumentation-langchain/test/instrumentationUtils.test.ts

@@ -1,11 +1,11 @@
-import { Tracer } from "@opentelemetry/api";
 import { addTracerToHandlers } from "../src/instrumentationUtils";
 import { LangChainTracer } from "../src/tracer";
 import { CallbackManager } from "@langchain/core/callbacks/manager";
+import { OITracer } from "@arizeai/openinference-core";
 
 describe("addTracerToHandlers", () => {
   it("should add a tracer if there are no handlers", () => {
-    const tracer = {} as Tracer;
+    const tracer = {} as OITracer;
 
     const result = addTracerToHandlers(tracer);
 
@@ -16,7 +16,7 @@ describe("addTracerToHandlers", () => {
     }
   });
   it("should add a handler to a pre-existing array of handlers", () => {
-    const tracer = {} as Tracer;
+    const tracer = {} as OITracer;
     const handlers = [new CallbackManager()];
 
     const result = addTracerToHandlers(tracer, handlers);
@@ -28,7 +28,7 @@ describe("addTracerToHandlers", () => {
     }
   });
   it("should add a handler to a callback handler class' handlers", () => {
-    const tracer = {} as Tracer;
+    const tracer = {} as OITracer;
     const handlers = new CallbackManager();
 
     const result = addTracerToHandlers(tracer, handlers);
@@ -43,7 +43,7 @@ describe("addTracerToHandlers", () => {
   });
 
   it("should not add a handler if it already exists in an array of handlers", () => {
-    const tracer = {} as Tracer;
+    const tracer = {} as OITracer;
     const handlers = [new LangChainTracer(tracer)];
 
     const result = addTracerToHandlers(tracer, handlers);
@@ -56,7 +56,7 @@ describe("addTracerToHandlers", () => {
   });
 
   it("should not add a handler if it already exists in a callback handler class' handlers", () => {
-    const tracer = {} as Tracer;
+    const tracer = {} as OITracer;
     const handlers = new CallbackManager();
     handlers.addHandler(new LangChainTracer(tracer));