From c03a5b60206459f2d213fa909bb68f4e931a7ca0 Mon Sep 17 00:00:00 2001 From: Parker Stafford <52351508+Parker-Stafford@users.noreply.github.com> Date: Wed, 23 Oct 2024 15:39:23 -0700 Subject: [PATCH] feat: add OITracer to autoinstrumentations providing support for attribute masking and context attribute inheritance (#1078) Co-authored-by: Tony Powell --- js/.changeset/khaki-students-attend.md | 6 + js/DEVELOPMENT.md | 77 +++++++- .../test/trace-config/OITracer.test.ts | 52 +++++- .../src/instrumentation.ts | 36 +++- .../src/instrumentationUtils.ts | 8 +- .../src/tracer.ts | 6 +- .../test/instrumentationUtils.test.ts | 12 +- .../test/langchainV1.test.ts | 146 +++++++++++++++- .../test/langchainV2.test.ts | 144 ++++++++++++++- .../jest.config.js | 1 + .../src/instrumentation.ts | 64 +++++-- .../test/openai.test.ts | 164 +++++++++++++++++- js/packages/openinference-vercel/README.md | 7 +- 13 files changed, 657 insertions(+), 66 deletions(-) create mode 100644 js/.changeset/khaki-students-attend.md diff --git a/js/.changeset/khaki-students-attend.md b/js/.changeset/khaki-students-attend.md new file mode 100644 index 000000000..5a2c9dc31 --- /dev/null +++ b/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 diff --git a/js/DEVELOPMENT.md b/js/DEVELOPMENT.md index 3695b2309..f9c1c7a4f 100644 --- a/js/DEVELOPMENT.md +++ b/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 diff --git a/js/packages/openinference-core/test/trace-config/OITracer.test.ts b/js/packages/openinference-core/test/trace-config/OITracer.test.ts index 08b3f8437..b1cb4b8c3 100644 --- a/js/packages/openinference-core/test/trace-config/OITracer.test.ts +++ b/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; let mockSpan: jest.Mocked; 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; }); - + 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, + ); + }); }); }); diff --git a/js/packages/openinference-instrumentation-langchain/src/instrumentation.ts b/js/packages/openinference-instrumentation-langchain/src/instrumentation.ts index d60862ddb..e86ac7a31 100644 --- a/js/packages/openinference-instrumentation-langchain/src/instrumentation.ts +++ b/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 { - 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 = {}; - constructor(tracer: Tracer) { + constructor(tracer: OITracer) { super(); this.tracer = tracer; } diff --git a/js/packages/openinference-instrumentation-langchain/test/instrumentationUtils.test.ts b/js/packages/openinference-instrumentation-langchain/test/instrumentationUtils.test.ts index f20512afa..1c9e5d17f 100644 --- a/js/packages/openinference-instrumentation-langchain/test/instrumentationUtils.test.ts +++ b/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)); diff --git a/js/packages/openinference-instrumentation-langchain/test/langchainV1.test.ts b/js/packages/openinference-instrumentation-langchain/test/langchainV1.test.ts index 9ce84bee4..37c87258e 100644 --- a/js/packages/openinference-instrumentation-langchain/test/langchainV1.test.ts +++ b/js/packages/openinference-instrumentation-langchain/test/langchainV1.test.ts @@ -17,11 +17,18 @@ import { SemanticConventions, } from "@arizeai/openinference-semantic-conventions"; import { LangChainTracer } from "../src/tracer"; -import { trace } from "@opentelemetry/api"; +import { context, trace } from "@opentelemetry/api"; import { completionsResponse, functionCallResponse } from "./fixtures"; import { DynamicTool } from "@langchain/coreV0.1/tools"; import { createRetrievalChain } from "langchainV0.1/chains/retrieval"; import { createStuffDocumentsChain } from "langchainV0.1/chains/combine_documents"; +import { + OITracer, + setAttributes, + setSession, +} from "@arizeai/openinference-core"; + +const memoryExporter = new InMemorySpanExporter(); const { INPUT_VALUE, @@ -48,12 +55,6 @@ const { METADATA, } = SemanticConventions; -const tracerProvider = new NodeTracerProvider(); -tracerProvider.register(); - -const instrumentation = new LangChainInstrumentation(); -instrumentation.disable(); - jest.mock("@langchain/openaiV0.1", () => { const originalModule = jest.requireActual("@langchain/openaiV0.1"); class MockChatOpenAI extends originalModule.ChatOpenAI { @@ -154,7 +155,10 @@ const expectedSpanAttributes = { }; describe("LangChainInstrumentation", () => { - const memoryExporter = new InMemorySpanExporter(); + const tracerProvider = new NodeTracerProvider(); + tracerProvider.register(); + const instrumentation = new LangChainInstrumentation(); + instrumentation.disable(); const provider = new NodeTracerProvider(); provider.getTracer("default"); @@ -524,6 +528,129 @@ describe("LangChainInstrumentation", () => { metadata: "{}", }); }); + + it("should capture context attributes and add them to spans", async () => { + await context.with( + setSession( + setAttributes(context.active(), { + "test-attribute": "test-value", + }), + { sessionId: "session-id" }, + ), + async () => { + const chatModel = new ChatOpenAI({ + openAIApiKey: "my-api-key", + modelName: "gpt-3.5-turbo", + temperature: 0, + }); + await chatModel.invoke("hello, this is a test"); + }, + ); + + const spans = memoryExporter.getFinishedSpans(); + expect(spans.length).toBe(1); + const span = spans[0]; + expect(span.attributes).toMatchInlineSnapshot(` +{ + "input.mime_type": "application/json", + "input.value": "{"messages":[[{"lc":1,"type":"constructor","id":["langchain_core","messages","HumanMessage"],"kwargs":{"content":"hello, this is a test","additional_kwargs":{},"response_metadata":{}}}]]}", + "llm.input_messages.0.message.content": "hello, this is a test", + "llm.input_messages.0.message.role": "user", + "llm.invocation_parameters": "{"model":"gpt-3.5-turbo","temperature":0,"top_p":1,"frequency_penalty":0,"presence_penalty":0,"n":1,"stream":false}", + "llm.model_name": "gpt-3.5-turbo", + "llm.output_messages.0.message.content": "This is a test.", + "llm.output_messages.0.message.role": "assistant", + "llm.token_count.completion": 5, + "llm.token_count.prompt": 12, + "llm.token_count.total": 17, + "metadata": "{}", + "openinference.span.kind": "LLM", + "output.mime_type": "application/json", + "output.value": "{"generations":[[{"text":"This is a test.","message":{"lc":1,"type":"constructor","id":["langchain_core","messages","AIMessage"],"kwargs":{"content":"This is a test.","tool_calls":[],"invalid_tool_calls":[],"additional_kwargs":{},"response_metadata":{"tokenUsage":{"completionTokens":5,"promptTokens":12,"totalTokens":17},"finish_reason":"stop"}}},"generationInfo":{"finish_reason":"stop"}}]],"llmOutput":{"tokenUsage":{"completionTokens":5,"promptTokens":12,"totalTokens":17}}}", + "session.id": "session-id", + "test-attribute": "test-value", +} +`); + }); +}); + +describe("LangChainInstrumentation with TraceConfigOptions", () => { + const tracerProvider = new NodeTracerProvider(); + tracerProvider.register(); + const instrumentation = new LangChainInstrumentation({ + traceConfig: { + hideInputs: true, + }, + }); + instrumentation.disable(); + const provider = new NodeTracerProvider(); + provider.getTracer("default"); + + instrumentation.setTracerProvider(tracerProvider); + tracerProvider.addSpanProcessor(new SimpleSpanProcessor(memoryExporter)); + + // @ts-expect-error the moduleExports property is private. This is needed to make the test work with auto-mocking + instrumentation._modules[0].moduleExports = CallbackManager; + beforeAll(() => { + instrumentation.enable(); + }); + afterAll(() => { + instrumentation.disable(); + }); + beforeEach(() => { + memoryExporter.reset(); + }); + afterEach(() => { + jest.resetAllMocks(); + jest.clearAllMocks(); + }); + it("should patch the callback manager module", async () => { + expect( + (CallbackManager as { openInferencePatched?: boolean }) + .openInferencePatched, + ).toBe(true); + }); + + it("should respect trace config options", async () => { + await context.with( + setSession( + setAttributes(context.active(), { + "test-attribute": "test-value", + }), + { sessionId: "session-id" }, + ), + async () => { + const chatModel = new ChatOpenAI({ + openAIApiKey: "my-api-key", + modelName: "gpt-3.5-turbo", + temperature: 0, + }); + await chatModel.invoke("hello, this is a test"); + }, + ); + + const spans = memoryExporter.getFinishedSpans(); + expect(spans.length).toBe(1); + const span = spans[0]; + expect(span.attributes).toMatchInlineSnapshot(` +{ + "input.value": "__REDACTED__", + "llm.invocation_parameters": "{"model":"gpt-3.5-turbo","temperature":0,"top_p":1,"frequency_penalty":0,"presence_penalty":0,"n":1,"stream":false}", + "llm.model_name": "gpt-3.5-turbo", + "llm.output_messages.0.message.content": "This is a test.", + "llm.output_messages.0.message.role": "assistant", + "llm.token_count.completion": 5, + "llm.token_count.prompt": 12, + "llm.token_count.total": 17, + "metadata": "{}", + "openinference.span.kind": "LLM", + "output.mime_type": "application/json", + "output.value": "{"generations":[[{"text":"This is a test.","message":{"lc":1,"type":"constructor","id":["langchain_core","messages","AIMessage"],"kwargs":{"content":"This is a test.","tool_calls":[],"invalid_tool_calls":[],"additional_kwargs":{},"response_metadata":{"tokenUsage":{"completionTokens":5,"promptTokens":12,"totalTokens":17},"finish_reason":"stop"}}},"generationInfo":{"finish_reason":"stop"}}]],"llmOutput":{"tokenUsage":{"completionTokens":5,"promptTokens":12,"totalTokens":17}}}", + "session.id": "session-id", + "test-attribute": "test-value", +} +`); + }); }); describe("LangChainTracer", () => { @@ -533,7 +660,8 @@ describe("LangChainTracer", () => { id: [], }; it("should delete runs after they are ended", async () => { - const langChainTracer = new LangChainTracer(trace.getTracer("default")); + const oiTracer = new OITracer({ tracer: trace.getTracer("default") }); + const langChainTracer = new LangChainTracer(oiTracer); for (let i = 0; i < 10; i++) { await langChainTracer.handleLLMStart(testSerialized, [], "runId"); expect(Object.keys(langChainTracer["runs"]).length).toBe(1); diff --git a/js/packages/openinference-instrumentation-langchain/test/langchainV2.test.ts b/js/packages/openinference-instrumentation-langchain/test/langchainV2.test.ts index cc1e02e32..1297a2461 100644 --- a/js/packages/openinference-instrumentation-langchain/test/langchainV2.test.ts +++ b/js/packages/openinference-instrumentation-langchain/test/langchainV2.test.ts @@ -23,8 +23,16 @@ import { LangChainTracer } from "../src/tracer"; import { trace } from "@opentelemetry/api"; import { completionsResponse, functionCallResponse } from "./fixtures"; import { DynamicTool } from "@langchain/core/tools"; +import { + OITracer, + setAttributes, + setSession, +} from "@arizeai/openinference-core"; +import { context } from "@opentelemetry/api"; jest.useFakeTimers(); +const memoryExporter = new InMemorySpanExporter(); + const { INPUT_VALUE, LLM_INPUT_MESSAGES, @@ -49,12 +57,6 @@ const { RETRIEVAL_DOCUMENTS, } = SemanticConventions; -const tracerProvider = new NodeTracerProvider(); -tracerProvider.register(); - -const instrumentation = new LangChainInstrumentation(); -instrumentation.disable(); - jest.mock("@langchain/openai", () => { const originalModule = jest.requireActual("@langchain/openai"); class MockChatOpenAI extends originalModule.ChatOpenAI { @@ -155,7 +157,10 @@ const expectedSpanAttributes = { }; describe("LangChainInstrumentation", () => { - const memoryExporter = new InMemorySpanExporter(); + const tracerProvider = new NodeTracerProvider(); + tracerProvider.register(); + const instrumentation = new LangChainInstrumentation(); + instrumentation.disable(); const provider = new NodeTracerProvider(); provider.getTracer("default"); @@ -535,6 +540,128 @@ describe("LangChainInstrumentation", () => { metadata: "{}", }); }); + + it("should capture context attributes and add them to spans", async () => { + await context.with( + setSession( + setAttributes(context.active(), { + "test-attribute": "test-value", + }), + { sessionId: "session-id" }, + ), + async () => { + const chatModel = new ChatOpenAI({ + openAIApiKey: "my-api-key", + modelName: "gpt-3.5-turbo", + temperature: 0, + }); + await chatModel.invoke("hello, this is a test"); + }, + ); + + const spans = memoryExporter.getFinishedSpans(); + expect(spans.length).toBe(1); + const span = spans[0]; + expect(span.attributes).toMatchInlineSnapshot(` +{ + "input.mime_type": "application/json", + "input.value": "{"messages":[[{"lc":1,"type":"constructor","id":["langchain_core","messages","HumanMessage"],"kwargs":{"content":"hello, this is a test","additional_kwargs":{},"response_metadata":{}}}]]}", + "llm.input_messages.0.message.content": "hello, this is a test", + "llm.input_messages.0.message.role": "user", + "llm.invocation_parameters": "{"model":"gpt-3.5-turbo","temperature":0,"top_p":1,"frequency_penalty":0,"presence_penalty":0,"n":1,"stream":false}", + "llm.model_name": "gpt-3.5-turbo", + "llm.output_messages.0.message.content": "This is a test.", + "llm.output_messages.0.message.role": "assistant", + "llm.token_count.completion": 5, + "llm.token_count.prompt": 12, + "llm.token_count.total": 17, + "metadata": "{"ls_provider":"openai","ls_model_name":"gpt-3.5-turbo","ls_model_type":"chat","ls_temperature":0}", + "openinference.span.kind": "LLM", + "output.mime_type": "application/json", + "output.value": "{"generations":[[{"text":"This is a test.","message":{"lc":1,"type":"constructor","id":["langchain_core","messages","AIMessage"],"kwargs":{"content":"This is a test.","tool_calls":[],"invalid_tool_calls":[],"additional_kwargs":{},"response_metadata":{"tokenUsage":{"completionTokens":5,"promptTokens":12,"totalTokens":17},"finish_reason":"stop"},"id":"chatcmpl-8adq9JloOzNZ9TyuzrKyLpGXexh6p"}},"generationInfo":{"finish_reason":"stop"}}]],"llmOutput":{"tokenUsage":{"completionTokens":5,"promptTokens":12,"totalTokens":17}}}", + "session.id": "session-id", + "test-attribute": "test-value", +} +`); + }); +}); + +describe("LangChainInstrumentation with TraceConfigOptions", () => { + const tracerProvider = new NodeTracerProvider(); + tracerProvider.register(); + const instrumentation = new LangChainInstrumentation({ + traceConfig: { + hideInputs: true, + }, + }); + instrumentation.disable(); + const provider = new NodeTracerProvider(); + provider.getTracer("default"); + instrumentation.setTracerProvider(tracerProvider); + tracerProvider.addSpanProcessor(new SimpleSpanProcessor(memoryExporter)); + + // @ts-expect-error the moduleExports property is private. This is needed to make the test work with auto-mocking + instrumentation._modules[0].moduleExports = CallbackManager; + beforeAll(() => { + instrumentation.enable(); + }); + afterAll(() => { + instrumentation.disable(); + }); + beforeEach(() => { + memoryExporter.reset(); + }); + afterEach(() => { + jest.resetAllMocks(); + jest.clearAllMocks(); + }); + it("should patch the callback manager module", async () => { + expect( + (CallbackManager as { openInferencePatched?: boolean }) + .openInferencePatched, + ).toBe(true); + }); + + it("should respect trace config options", async () => { + await context.with( + setSession( + setAttributes(context.active(), { + "test-attribute": "test-value", + }), + { sessionId: "session-id" }, + ), + async () => { + const chatModel = new ChatOpenAI({ + openAIApiKey: "my-api-key", + modelName: "gpt-3.5-turbo", + temperature: 0, + }); + await chatModel.invoke("hello, this is a test"); + }, + ); + + const spans = memoryExporter.getFinishedSpans(); + expect(spans.length).toBe(1); + const span = spans[0]; + expect(span.attributes).toMatchInlineSnapshot(` +{ + "input.value": "__REDACTED__", + "llm.invocation_parameters": "{"model":"gpt-3.5-turbo","temperature":0,"top_p":1,"frequency_penalty":0,"presence_penalty":0,"n":1,"stream":false}", + "llm.model_name": "gpt-3.5-turbo", + "llm.output_messages.0.message.content": "This is a test.", + "llm.output_messages.0.message.role": "assistant", + "llm.token_count.completion": 5, + "llm.token_count.prompt": 12, + "llm.token_count.total": 17, + "metadata": "{"ls_provider":"openai","ls_model_name":"gpt-3.5-turbo","ls_model_type":"chat","ls_temperature":0}", + "openinference.span.kind": "LLM", + "output.mime_type": "application/json", + "output.value": "{"generations":[[{"text":"This is a test.","message":{"lc":1,"type":"constructor","id":["langchain_core","messages","AIMessage"],"kwargs":{"content":"This is a test.","tool_calls":[],"invalid_tool_calls":[],"additional_kwargs":{},"response_metadata":{"tokenUsage":{"completionTokens":5,"promptTokens":12,"totalTokens":17},"finish_reason":"stop"},"id":"chatcmpl-8adq9JloOzNZ9TyuzrKyLpGXexh6p"}},"generationInfo":{"finish_reason":"stop"}}]],"llmOutput":{"tokenUsage":{"completionTokens":5,"promptTokens":12,"totalTokens":17}}}", + "session.id": "session-id", + "test-attribute": "test-value", +} +`); + }); }); describe("LangChainTracer", () => { @@ -544,7 +671,8 @@ describe("LangChainTracer", () => { id: [], }; it("should delete runs after they are ended", async () => { - const langChainTracer = new LangChainTracer(trace.getTracer("default")); + const oiTracer = new OITracer({ tracer: trace.getTracer("default") }); + const langChainTracer = new LangChainTracer(oiTracer); for (let i = 0; i < 10; i++) { await langChainTracer.handleLLMStart(testSerialized, [], "runId"); expect(Object.keys(langChainTracer["runs"]).length).toBe(1); diff --git a/js/packages/openinference-instrumentation-openai/jest.config.js b/js/packages/openinference-instrumentation-openai/jest.config.js index 2fd4f78cf..feefeda86 100644 --- a/js/packages/openinference-instrumentation-openai/jest.config.js +++ b/js/packages/openinference-instrumentation-openai/jest.config.js @@ -3,4 +3,5 @@ module.exports = { preset: "ts-jest", testEnvironment: "node", prettierPath: null, + testMatch: ["/test/**/*.test.ts"], }; diff --git a/js/packages/openinference-instrumentation-openai/src/instrumentation.ts b/js/packages/openinference-instrumentation-openai/src/instrumentation.ts index 3c2aa76de..580b95f4b 100644 --- a/js/packages/openinference-instrumentation-openai/src/instrumentation.ts +++ b/js/packages/openinference-instrumentation-openai/src/instrumentation.ts @@ -40,7 +40,11 @@ import { import { assertUnreachable, isString } from "./typeUtils"; import { isTracingSuppressed } from "@opentelemetry/core"; -import { safelyJSONStringify } from "@arizeai/openinference-core"; +import { + OITracer, + safelyJSONStringify, + TraceConfigOptions, +} from "@arizeai/openinference-core"; const MODULE_NAME = "openai"; @@ -74,13 +78,34 @@ function getExecContext(span: Span) { } return execContext; } +/** + * An auto instrumentation class for OpenAI that creates {@link https://github.com/Arize-ai/openinference/blob/main/spec/semantic_conventions.md|OpenInference} Compliant spans for the OpenAI API + * @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 OpenAIInstrumentation extends InstrumentationBase { - 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-openai", VERSION, - Object.assign({}, config), + Object.assign({}, instrumentationConfig), ); + this.oiTracer = new OITracer({ tracer: this.tracer, traceConfig }); } protected init(): InstrumentationModuleDefinition { @@ -131,7 +156,7 @@ export class OpenAIInstrumentation extends InstrumentationBase { ) { const body = args[0]; const { messages: _messages, ...invocationParameters } = body; - const span = instrumentation.tracer.startSpan( + const span = instrumentation.oiTracer.startSpan( `OpenAI Chat Completions`, { kind: SpanKind.INTERNAL, @@ -216,19 +241,22 @@ export class OpenAIInstrumentation extends InstrumentationBase { ) { const body = args[0]; const { prompt: _prompt, ...invocationParameters } = body; - const span = instrumentation.tracer.startSpan(`OpenAI Completions`, { - kind: SpanKind.INTERNAL, - attributes: { - [SemanticConventions.OPENINFERENCE_SPAN_KIND]: - OpenInferenceSpanKind.LLM, - [SemanticConventions.LLM_MODEL_NAME]: body.model, - [SemanticConventions.LLM_INVOCATION_PARAMETERS]: - JSON.stringify(invocationParameters), - [SemanticConventions.LLM_SYSTEM]: LLMSystem.OPENAI, - [SemanticConventions.LLM_PROVIDER]: LLMProvider.OPENAI, - ...getCompletionInputValueAndMimeType(body), + const span = instrumentation.oiTracer.startSpan( + `OpenAI Completions`, + { + kind: SpanKind.INTERNAL, + attributes: { + [SemanticConventions.OPENINFERENCE_SPAN_KIND]: + OpenInferenceSpanKind.LLM, + [SemanticConventions.LLM_MODEL_NAME]: body.model, + [SemanticConventions.LLM_INVOCATION_PARAMETERS]: + JSON.stringify(invocationParameters), + [SemanticConventions.LLM_SYSTEM]: LLMSystem.OPENAI, + [SemanticConventions.LLM_PROVIDER]: LLMProvider.OPENAI, + ...getCompletionInputValueAndMimeType(body), + }, }, - }); + ); const execContext = getExecContext(span); const execPromise = safeExecuteInTheMiddle< @@ -287,7 +315,7 @@ export class OpenAIInstrumentation extends InstrumentationBase { const body = args[0]; const { input } = body; const isStringInput = typeof input === "string"; - const span = instrumentation.tracer.startSpan(`OpenAI Embeddings`, { + const span = instrumentation.oiTracer.startSpan(`OpenAI Embeddings`, { kind: SpanKind.INTERNAL, attributes: { [SemanticConventions.OPENINFERENCE_SPAN_KIND]: @@ -345,7 +373,7 @@ export class OpenAIInstrumentation extends InstrumentationBase { // 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; diff --git a/js/packages/openinference-instrumentation-openai/test/openai.test.ts b/js/packages/openinference-instrumentation-openai/test/openai.test.ts index 4f636b49b..c39630336 100644 --- a/js/packages/openinference-instrumentation-openai/test/openai.test.ts +++ b/js/packages/openinference-instrumentation-openai/test/openai.test.ts @@ -6,14 +6,10 @@ import { import { NodeTracerProvider } from "@opentelemetry/sdk-trace-node"; import { suppressTracing } from "@opentelemetry/core"; import { context } from "@opentelemetry/api"; -const tracerProvider = new NodeTracerProvider(); -tracerProvider.register(); - -const instrumentation = new OpenAIInstrumentation(); -instrumentation.disable(); import * as OpenAI from "openai"; import { Stream } from "openai/streaming"; +import { setPromptTemplate, setSession } from "@arizeai/openinference-core"; // Function tools async function getCurrentLocation() { @@ -24,11 +20,15 @@ async function getWeather(_args: { location: string }) { return { temperature: 52, precipitation: "rainy" }; } +const memoryExporter = new InMemorySpanExporter(); + describe("OpenAIInstrumentation", () => { + const tracerProvider = new NodeTracerProvider(); + tracerProvider.register(); + const instrumentation = new OpenAIInstrumentation(); + instrumentation.disable(); let openai: OpenAI.OpenAI; - const memoryExporter = new InMemorySpanExporter(); - instrumentation.setTracerProvider(tracerProvider); tracerProvider.addSpanProcessor(new SimpleSpanProcessor(memoryExporter)); // @ts-expect-error the moduleExports property is private. This is needed to make the test work with auto-mocking @@ -838,6 +838,156 @@ describe("OpenAIInstrumentation", () => { "output.mime_type": "application/json", "output.value": "{"id":"chatcmpl-8adq9JloOzNZ9TyuzrKyLpGXexh6p","object":"chat.completion","created":1703743645,"model":"gpt-3.5-turbo-0613","choices":[{"index":0,"message":{"role":"assistant","content":"This is a test."},"logprobs":null,"finish_reason":"stop"}],"usage":{"prompt_tokens":12,"completion_tokens":5,"total_tokens":17}}", } +`); + }); + + it("should capture context attributes and add them to spans", async () => { + const response = { + id: "cmpl-8fZu1H3VijJUWev9asnxaYyQvJTC9", + object: "text_completion", + created: 1704920149, + model: "gpt-3.5-turbo-instruct", + choices: [ + { + text: "This is a test", + index: 0, + logprobs: null, + finish_reason: "stop", + }, + ], + usage: { prompt_tokens: 12, completion_tokens: 5, total_tokens: 17 }, + }; + // Mock out the completions endpoint + jest.spyOn(openai, "post").mockImplementation( + // @ts-expect-error the response type is not correct - this is just for testing + async (): Promise => { + return response; + }, + ); + await context.with( + setSession( + setPromptTemplate(context.active(), { + template: "hello {name}", + variables: { name: "world" }, + version: "V1.0", + }), + { sessionId: "session-id" }, + ), + async () => { + await openai.completions.create({ + prompt: "Say this is a test", + model: "gpt-3.5-turbo-instruct", + }); + }, + ); + const spans = memoryExporter.getFinishedSpans(); + expect(spans.length).toBe(1); + const span = spans[0]; + expect(span.name).toBe("OpenAI Completions"); + expect(span.attributes).toMatchInlineSnapshot(` +{ + "input.mime_type": "text/plain", + "input.value": "Say this is a test", + "llm.invocation_parameters": "{"model":"gpt-3.5-turbo-instruct"}", + "llm.model_name": "gpt-3.5-turbo-instruct", + "llm.prompt_template.template": "hello {name}", + "llm.prompt_template.variables": "{"name":"world"}", + "llm.prompt_template.version": "V1.0", + "llm.provider": "openai", + "llm.system": "openai", + "llm.token_count.completion": 5, + "llm.token_count.prompt": 12, + "llm.token_count.total": 17, + "openinference.span.kind": "LLM", + "output.mime_type": "text/plain", + "output.value": "This is a test", + "session.id": "session-id", +} +`); + }); +}); + +describe("OpenAIInstrumentation with TraceConfig", () => { + const tracerProvider = new NodeTracerProvider(); + tracerProvider.register(); + const instrumentation = new OpenAIInstrumentation({ + traceConfig: { hideInputs: true }, + }); + instrumentation.disable(); + let openai: OpenAI.OpenAI; + + instrumentation.setTracerProvider(tracerProvider); + tracerProvider.addSpanProcessor(new SimpleSpanProcessor(memoryExporter)); + // @ts-expect-error the moduleExports property is private. This is needed to make the test work with auto-mocking + instrumentation._modules[0].moduleExports = OpenAI; + + beforeAll(() => { + instrumentation.enable(); + openai = new OpenAI.OpenAI({ + apiKey: "fake-api-key", + }); + }); + afterAll(() => { + instrumentation.disable(); + }); + beforeEach(() => { + memoryExporter.reset(); + }); + afterEach(() => { + jest.clearAllMocks(); + }); + it("is patched", () => { + expect( + (OpenAI as { openInferencePatched?: boolean }).openInferencePatched, + ).toBe(true); + expect(isPatched()).toBe(true); + }); + it("should respect a trace config and mask attributes accordingly", async () => { + const response = { + id: "cmpl-8fZu1H3VijJUWev9asnxaYyQvJTC9", + object: "text_completion", + created: 1704920149, + model: "gpt-3.5-turbo-instruct", + choices: [ + { + text: "This is a test", + index: 0, + logprobs: null, + finish_reason: "stop", + }, + ], + usage: { prompt_tokens: 12, completion_tokens: 5, total_tokens: 17 }, + }; + // Mock out the completions endpoint + jest.spyOn(openai, "post").mockImplementation( + // @ts-expect-error the response type is not correct - this is just for testing + async (): Promise => { + return response; + }, + ); + + await openai.completions.create({ + prompt: "Say this is a test", + model: "gpt-3.5-turbo-instruct", + }); + const spans = memoryExporter.getFinishedSpans(); + expect(spans.length).toBe(1); + const span = spans[0]; + expect(span.name).toBe("OpenAI Completions"); + expect(span.attributes).toMatchInlineSnapshot(` +{ + "input.value": "__REDACTED__", + "llm.invocation_parameters": "{"model":"gpt-3.5-turbo-instruct"}", + "llm.model_name": "gpt-3.5-turbo-instruct", + "llm.provider": "openai", + "llm.system": "openai", + "llm.token_count.completion": 5, + "llm.token_count.prompt": 12, + "llm.token_count.total": 17, + "openinference.span.kind": "LLM", + "output.mime_type": "text/plain", + "output.value": "This is a test", +} `); }); }); diff --git a/js/packages/openinference-vercel/README.md b/js/packages/openinference-vercel/README.md index a24c7c9e7..efba1dc17 100644 --- a/js/packages/openinference-vercel/README.md +++ b/js/packages/openinference-vercel/README.md @@ -24,7 +24,8 @@ npm i @opentelemetry/api @vercel/otel @opentelemetry/exporter-trace-otlp-proto @ To process your Vercel AI SDK Spans add a `OpenInferenceSimpleSpanProcessor` or `OpenInferenceBatchSpanProcessor` to your OpenTelemetry configuration. -> Note: The `OpenInferenceSpanProcessor` does not handle the exporting of spans so you will pass it an [exporter](https://opentelemetry.io/docs/languages/js/exporters/) as a parameter. +> [!NOTE] +> The `OpenInferenceSpanProcessor` does not handle the exporting of spans so you will pass it an [exporter](https://opentelemetry.io/docs/languages/js/exporters/) as a parameter. ```typescript import { registerOTel } from "@vercel/otel"; @@ -43,14 +44,14 @@ export function register() { registerOTel({ serviceName: "phoenix-next-app", attributes: { - // This is not required but it will + // This is not required but it will ensure your traces get added to a specific project in Arize Phoenix [SEMRESATTRS_PROJECT_NAME]: "your-next-app", }, spanProcessors: [ new OpenInferenceSimpleSpanProcessor({ exporter: new OTLPTraceExporter({ headers: { - // API key if you are sending it ot Phoenix + // API key if you are sending it to Phoenix api_key: process.env["PHOENIX_API_KEY"], }, url: