Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 21 additions & 0 deletions e2e/scenarios/openai-instrumentation/assertions.ts
Original file line number Diff line number Diff line change
Expand Up @@ -427,6 +427,7 @@ function buildPayloadSummary(events: CapturedLogEvent[]): Json {
}

export function defineOpenAIInstrumentationAssertions(options: {
assertPrivateFieldMethodsOperation?: boolean;
name: string;
runScenario: RunOpenAIScenario;
snapshotName: string;
Expand Down Expand Up @@ -466,6 +467,26 @@ export function defineOpenAIInstrumentationAssertions(options: {
});
});

if (options.assertPrivateFieldMethodsOperation) {
test(
"keeps wrapped v6 client private-field methods callable",
testConfig,
() => {
const root = findLatestSpan(events, ROOT_NAME);
const operation = findLatestSpan(
events,
"openai-client-private-fields-operation",
);

expect(operation).toBeDefined();
expect(operation?.row.metadata).toMatchObject({
operation: "client-private-fields",
});
expect(operation?.span.parentIds).toEqual([root?.span.id ?? ""]);
},
);
}

for (const spec of OPERATION_SPECS) {
test(spec.testName, testConfig, () => {
const root = findLatestSpan(events, ROOT_NAME);
Expand Down
51 changes: 51 additions & 0 deletions e2e/scenarios/openai-instrumentation/scenario.impl.mjs
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,15 @@ async function awaitMaybeWithResponse(request) {
};
}

function parseMajorVersion(version) {
if (typeof version !== "string") {
return null;
}

const major = Number.parseInt(version.split(".")[0], 10);
return Number.isNaN(major) ? null : major;
}

export async function runOpenAIInstrumentationScenario(options) {
const baseClient = new options.OpenAI({
apiKey: process.env.OPENAI_API_KEY,
Expand All @@ -52,9 +61,51 @@ export async function runOpenAIInstrumentationScenario(options) {
const client = options.decorateClient
? options.decorateClient(baseClient)
: baseClient;
const openAIMajorVersion = parseMajorVersion(options.openaiSdkVersion);
const shouldCheckPrivateFieldMethods =
typeof options.decorateClient === "function" &&
openAIMajorVersion !== null &&
openAIMajorVersion >= 6;

await runTracedScenario({
callback: async () => {
if (shouldCheckPrivateFieldMethods) {
await runOperation(
"openai-client-private-fields-operation",
"client-private-fields",
async () => {
if (
typeof client.buildURL !== "function" ||
typeof client.buildRequest !== "function"
) {
throw new Error(
"Expected wrapped OpenAI v6 client to expose buildURL and buildRequest",
);
}

const builtUrl = client.buildURL("/files", null);
if (typeof builtUrl !== "string" || !builtUrl.includes("/files")) {
throw new Error(
`Unexpected buildURL result: ${String(builtUrl)}`,
);
}

const builtRequest = await client.buildRequest(
{ method: "post", path: "/files" },
{ retryCount: 0 },
);
if (
typeof builtRequest?.url !== "string" ||
!builtRequest.url.includes("/files")
) {
throw new Error(
`Unexpected buildRequest result: ${String(builtRequest?.url)}`,
);
}
},
);
}

await runOperation("openai-chat-operation", "chat", async () => {
await client.chat.completions.create({
model: OPENAI_MODEL,
Expand Down
6 changes: 6 additions & 0 deletions e2e/scenarios/openai-instrumentation/scenario.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -14,12 +14,14 @@ const openaiScenarios = await Promise.all(
[
{
autoEntry: "scenario.openai-v4.mjs",
disablePrivateFieldMethodsAssertion: true,
dependencyName: "openai-v4",
snapshotName: "openai-v4",
wrapperEntry: "scenario.openai-v4.ts",
},
{
autoEntry: "scenario.openai-v5.mjs",
disablePrivateFieldMethodsAssertion: true,
dependencyName: "openai-v5",
snapshotName: "openai-v5",
wrapperEntry: "scenario.openai-v5.ts",
Expand All @@ -40,8 +42,12 @@ const openaiScenarios = await Promise.all(
);

for (const scenario of openaiScenarios) {
const assertPrivateFieldMethodsOperation =
!scenario.disablePrivateFieldMethodsAssertion;

describe(`openai sdk ${scenario.version}`, () => {
defineOpenAIInstrumentationAssertions({
assertPrivateFieldMethodsOperation,
name: "wrapped instrumentation",
runScenario: async ({ runScenarioDir }) => {
await runScenarioDir({
Expand Down
14 changes: 14 additions & 0 deletions js/src/wrappers/oai.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,20 @@ test("openai is installed", () => {
assert.ok(OpenAI);
});

test("wrapOpenAI keeps private-field helper methods callable", async () => {
const wrapped = wrapOpenAI(new OpenAI({ apiKey: "test-key" }));

expect(wrapped.buildURL).toBe(wrapped.buildURL);
expect(wrapped.buildRequest).toBe(wrapped.buildRequest);
expect(() => wrapped.buildURL("/files", null)).not.toThrow();

const request = await wrapped.buildRequest(
{ method: "post", path: "/files" },
{ retryCount: 0 },
);
expect(request.url).toContain("/files");
});

describe("openai client unit tests", TEST_SUITE_OPTIONS, () => {
let oai: OpenAI;
let client: OpenAI;
Expand Down
49 changes: 45 additions & 4 deletions js/src/wrappers/oai.ts
Original file line number Diff line number Diff line change
Expand Up @@ -74,6 +74,12 @@ type OpenAILike = OpenAIV4Client;
export function wrapOpenAIv4<T extends OpenAILike>(openai: T): T {
// eslint-disable-next-line @typescript-eslint/consistent-type-assertions
const typedOpenai = openai as OpenAIV4Client;
// Recover `this` for fallback methods so private fields and internal slots
// keep seeing the original OpenAI instance instead of the proxy.
const privateMethodWorkaroundCache = new WeakMap<
(...args: unknown[]) => unknown,
(...args: unknown[]) => unknown
>();

const completionProxy = new Proxy(typedOpenai.chat.completions, {
get(target, name, receiver) {
Expand Down Expand Up @@ -142,8 +148,8 @@ export function wrapOpenAIv4<T extends OpenAILike>(openai: T): T {
}

// eslint-disable-next-line @typescript-eslint/consistent-type-assertions
return new Proxy(typedOpenai, {
get(target, name, receiver) {
const topLevelProxy = new Proxy(typedOpenai, {
get(target, name) {
switch (name) {
case "chat":
return chatProxy;
Expand All @@ -158,9 +164,44 @@ export function wrapOpenAIv4<T extends OpenAILike>(openai: T): T {
if (name === "beta" && betaProxy) {
return betaProxy;
}
return Reflect.get(target, name, receiver);

// The following rather convoluted code is a workaround for https://github.com/braintrustdata/braintrust-sdk-javascript/issues/1693
// The problem is that Proxies are inherently difficult to work with native private class fields because when a
// class function accesses a private field, JS checks whether `this` is equal to the actual instance with the
// private field and if that's not the case, it throws a `TypeError`.
// We could have also done `if (typeof value === "function") return value.bind(target);`, but it would have
// created a new function on each function access, so we are caching, and also it would have always stomped on
// someone passing another `this` which may clash with different instrumentations.

// Use the real client as receiver when reading fallback members.
const value = Reflect.get(target, name, target);
if (typeof value !== "function") {
return value;
}

const cachedValue = privateMethodWorkaroundCache.get(value);
if (cachedValue) {
return cachedValue;
}

const thisBoundValue = function (
this: unknown,
...args: unknown[]
): unknown {
// Calling through the proxy would set `this` to the proxy and break
// native private-field methods, so recover the original target.
const thisArg = this === topLevelProxy ? target : this;
const output = Reflect.apply(value, thisArg, args);
// Preserve chaining on wrapped clients (method returns `this`).
return output === target ? topLevelProxy : output;
};

privateMethodWorkaroundCache.set(value, thisBoundValue);
return thisBoundValue;
},
}) as T;
});

return topLevelProxy as T;
}

type SpanInfo = {
Expand Down
Loading