Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix: log instrumentation patching #148

Merged
merged 3 commits into from
Mar 15, 2024
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
17 changes: 13 additions & 4 deletions packages/instrumentation-azure/src/instrumentation.ts
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,8 @@ export class AzureOpenAIInstrumentation extends InstrumentationBase<any> {
}

public manuallyInstrument(module: typeof azure) {
this._diag.debug(`Patching @azure/openai manually`);

this._wrap(
module.OpenAIClient.prototype,
"getChatCompletions",
Expand All @@ -74,7 +76,9 @@ export class AzureOpenAIInstrumentation extends InstrumentationBase<any> {
return module;
}

private patch(moduleExports: typeof azure) {
private patch(moduleExports: typeof azure, moduleVersion?: string) {
this._diag.debug(`Patching @azure/openai@${moduleVersion}`);

this._wrap(
moduleExports.OpenAIClient.prototype,
"getChatCompletions",
Expand All @@ -88,7 +92,9 @@ export class AzureOpenAIInstrumentation extends InstrumentationBase<any> {
return moduleExports;
}

private unpatch(moduleExports: typeof azure): void {
private unpatch(moduleExports: typeof azure, moduleVersion?: string): void {
this._diag.debug(`Unpatching @azure/openai@${moduleVersion}`);

this._unwrap(moduleExports.OpenAIClient.prototype, "getChatCompletions");
this._unwrap(moduleExports.OpenAIClient.prototype, "getCompletions");
}
Expand Down Expand Up @@ -127,8 +133,11 @@ export class AzureOpenAIInstrumentation extends InstrumentationBase<any> {
return original.apply(this, args);
});
},
// eslint-disable-next-line @typescript-eslint/no-empty-function
() => {},
(e) => {
if (e) {
plugin._diag.error("Error in Azure OpenAI instrumentation", e);
}
},
);

const wrappedPromise = plugin._wrapPromise(
Expand Down
21 changes: 17 additions & 4 deletions packages/instrumentation-bedrock/src/instrumentation.ts
Original file line number Diff line number Diff line change
Expand Up @@ -58,14 +58,20 @@ export class BedrockInstrumentation extends InstrumentationBase<any> {
}

public manuallyInstrument(module: typeof bedrock) {
this._diag.debug(`Patching @aws-sdk/client-bedrock-runtime manually`);

this._wrap(
module.BedrockRuntimeClient.prototype,
"send",
this.wrapperMethod(),
);
}

private wrap(module: typeof bedrock) {
private wrap(module: typeof bedrock, moduleVersion?: string) {
this._diag.debug(
`Patching @aws-sdk/client-bedrock-runtime@${moduleVersion}`,
);

this._wrap(
module.BedrockRuntimeClient.prototype,
"send",
Expand All @@ -75,7 +81,11 @@ export class BedrockInstrumentation extends InstrumentationBase<any> {
return module;
}

private unwrap(module: typeof bedrock) {
private unwrap(module: typeof bedrock, moduleVersion?: string) {
this._diag.debug(
`Unpatching @aws-sdk/client-bedrock-runtime@${moduleVersion}`,
);

this._unwrap(module.BedrockRuntimeClient.prototype, "send");
}

Expand All @@ -95,8 +105,11 @@ export class BedrockInstrumentation extends InstrumentationBase<any> {
return original.apply(this, args);
});
},
// eslint-disable-next-line @typescript-eslint/no-empty-function
() => {},
(e) => {
if (e) {
plugin._diag.error(`Error in bedrock instrumentation`, e);
}
},
);
const wrappedPromise = plugin._wrapPromise(span, execPromise);
return context.bind(execContext, wrappedPromise);
Expand Down
16 changes: 12 additions & 4 deletions packages/instrumentation-cohere/src/instrumentation.ts
Original file line number Diff line number Diff line change
Expand Up @@ -60,10 +60,13 @@ export class CohereInstrumentation extends InstrumentationBase<any> {
}

public manuallyInstrument(module: typeof cohere) {
this._diag.debug(`Manually patching cohere-ai`);
this.wrap(module);
}

private wrap(module: typeof cohere) {
private wrap(module: typeof cohere, moduleVersion?: string) {
this._diag.debug(`Patching cohere-ai@${moduleVersion}`);

this._wrap(
module.CohereClient.prototype,
"generate",
Expand Down Expand Up @@ -93,7 +96,9 @@ export class CohereInstrumentation extends InstrumentationBase<any> {
return module;
}

private unwrap(module: typeof cohere) {
private unwrap(module: typeof cohere, moduleVersion?: string) {
this._diag.debug(`Unpatching @cohere-ai@${moduleVersion}`);

this._unwrap(module.CohereClient.prototype, "generateStream");
this._unwrap(module.CohereClient.prototype, "chat");
this._unwrap(module.CohereClient.prototype, "chatStream");
Expand All @@ -117,8 +122,11 @@ export class CohereInstrumentation extends InstrumentationBase<any> {
return original.apply(this, args);
});
},
// eslint-disable-next-line @typescript-eslint/no-empty-function
() => {},
(e) => {
if (e) {
plugin._diag.error("Error in cohere instrumentation", e);
}
},
);
const wrappedPromise = plugin._wrapPromise(type, span, execPromise);
return context.bind(execContext, wrappedPromise as any);
Expand Down
44 changes: 37 additions & 7 deletions packages/instrumentation-langchain/src/instrumentation.ts
Original file line number Diff line number Diff line change
Expand Up @@ -42,12 +42,15 @@ export class LangChainInstrumentation extends InstrumentationBase<any> {
toolsModule?: any;
}) {
if (chainsModule) {
this._diag.debug("Manually instrumenting langchain chains");
this.patchChainModule(chainsModule);
}
if (agentsModule) {
this._diag.debug("Manually instrumenting langchain agents");
this.patchAgentModule(agentsModule);
}
if (toolsModule) {
this._diag.debug("Manually instrumenting langchain tools");
this.patchToolsModule(toolsModule);
}
}
Expand All @@ -74,7 +77,12 @@ export class LangChainInstrumentation extends InstrumentationBase<any> {
return [chainModule, agentModule, toolsModule];
}

private patchChainModule(moduleExports: typeof ChainsModule) {
private patchChainModule(
moduleExports: typeof ChainsModule,
moduleVersion?: string,
) {
this._diag.debug(`Patching langchain/chains.cjs@${moduleVersion}`);

this._wrap(
moduleExports.RetrievalQAChain.prototype,
"_call",
Expand All @@ -92,7 +100,12 @@ export class LangChainInstrumentation extends InstrumentationBase<any> {
return moduleExports;
}

private patchAgentModule(moduleExports: typeof AgentsModule) {
private patchAgentModule(
moduleExports: typeof AgentsModule,
moduleVersion?: string,
) {
this._diag.debug(`Patching langchain/agents.cjs@${moduleVersion}`);

this._wrap(
moduleExports.AgentExecutor.prototype,
"_call",
Expand All @@ -105,7 +118,12 @@ export class LangChainInstrumentation extends InstrumentationBase<any> {
return moduleExports;
}

private patchToolsModule(moduleExports: typeof ToolsModule) {
private patchToolsModule(
moduleExports: typeof ToolsModule,
moduleVersion?: string,
) {
this._diag.debug(`Patching langchain/tools.cjs@${moduleVersion}`);

this._wrap(
moduleExports.Tool.prototype,
"call",
Expand All @@ -114,19 +132,31 @@ export class LangChainInstrumentation extends InstrumentationBase<any> {
return moduleExports;
}

private unpatchChainModule(moduleExports: any) {
private unpatchChainModule(
moduleExports: typeof ChainsModule,
moduleVersion?: string,
) {
this._diag.debug(`Unpatching langchain/chains.cjs@${moduleVersion}`);

this._unwrap(moduleExports.RetrievalQAChain.prototype, "_call");
this._unwrap(moduleExports.BaseChain.prototype, "call");
return moduleExports;
}

private unpatchAgentModule(moduleExports: any) {
private unpatchAgentModule(
moduleExports: typeof AgentsModule,
moduleVersion?: string,
) {
this._diag.debug(`Unpatching langchain/agents.cjs@${moduleVersion}`);

this._unwrap(moduleExports.AgentExecutor.prototype, "_call");
return moduleExports;
}

private unpatchToolsModule(moduleExports: any) {
this._unwrap(moduleExports.AgentExecutor.prototype, "_call");
private unpatchToolsModule(moduleExports: typeof ToolsModule) {
this._diag.debug(`Unpatching langchain/tools.cjs`);

this._unwrap(moduleExports.Tool.prototype, "call");
return moduleExports;
}

Expand Down
10 changes: 8 additions & 2 deletions packages/instrumentation-llamaindex/src/instrumentation.ts
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,8 @@ export class LlamaIndexInstrumentation extends InstrumentationBase<any> {
}

public manuallyInstrument(module: typeof llamaindex) {
this._diag.debug("Manually instrumenting llamaindex");

this.patch(module);
}

Expand Down Expand Up @@ -78,7 +80,9 @@ export class LlamaIndexInstrumentation extends InstrumentationBase<any> {
return retriever && (retriever as BaseRetriever).retrieve !== undefined;
}

private patch(moduleExports: typeof llamaindex) {
private patch(moduleExports: typeof llamaindex, moduleVersion?: string) {
this._diag.debug(`Patching llamaindex@${moduleVersion}`);

const customLLMInstrumentation = new CustomLLMInstrumentation(
this._config,
() => this.tracer, // this is on purpose. Tracer may change
Expand Down Expand Up @@ -158,7 +162,9 @@ export class LlamaIndexInstrumentation extends InstrumentationBase<any> {
return moduleExports;
}

private unpatch(moduleExports: typeof llamaindex) {
private unpatch(moduleExports: typeof llamaindex, moduleVersion?: string) {
this._diag.debug(`Unpatching llamaindex@${moduleVersion}`);

this._unwrap(moduleExports.RetrieverQueryEngine.prototype, "query");

for (const key in moduleExports) {
Expand Down
17 changes: 13 additions & 4 deletions packages/instrumentation-openai/src/instrumentation.ts
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,8 @@ export class OpenAIInstrumentation extends InstrumentationBase<any> {
}

public manuallyInstrument(module: typeof openai.OpenAI) {
this._diag.debug(`Manually instrumenting openai`);

// Old version of OpenAI API (v3.1.0)
if ((module as any).OpenAIApi) {
this._wrap(
Expand Down Expand Up @@ -92,7 +94,9 @@ export class OpenAIInstrumentation extends InstrumentationBase<any> {
return module;
}

private patch(moduleExports: typeof openai) {
private patch(moduleExports: typeof openai, moduleVersion?: string) {
this._diag.debug(`Patching openai@${moduleVersion}`);

// Old version of OpenAI API (v3.1.0)
if ((moduleExports as any).OpenAIApi) {
this._wrap(
Expand Down Expand Up @@ -120,7 +124,9 @@ export class OpenAIInstrumentation extends InstrumentationBase<any> {
return moduleExports;
}

private unpatch(moduleExports: typeof openai): void {
private unpatch(moduleExports: typeof openai, moduleVersion?: string): void {
this._diag.debug(`Unpatching openai@${moduleVersion}`);

// Old version of OpenAI API (v3.1.0)
if ((moduleExports as any).OpenAIApi) {
this._unwrap(
Expand Down Expand Up @@ -171,8 +177,11 @@ export class OpenAIInstrumentation extends InstrumentationBase<any> {
return original.apply(this, args);
});
},
// eslint-disable-next-line @typescript-eslint/no-empty-function
() => {},
(e) => {
if (e) {
plugin._diag.error("OpenAI instrumentation: error", e);
}
},
);

if (
Expand Down
29 changes: 23 additions & 6 deletions packages/instrumentation-pinecone/src/instrumentation.ts
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,9 @@ export class PineconeInstrumentation extends InstrumentationBase<any> {
return module;
}

private patch(moduleExports: typeof pinecone) {
private patch(moduleExports: typeof pinecone, moduleVersion?: string) {
this._diag.debug(`Patching @pinecone-database/pinecone@${moduleVersion}`);

this._wrap(
moduleExports.Index.prototype,
"query",
Expand Down Expand Up @@ -77,7 +79,12 @@ export class PineconeInstrumentation extends InstrumentationBase<any> {
return moduleExports;
}

private unpatch(moduleExports: typeof pinecone): void {
private unpatch(
moduleExports: typeof pinecone,
moduleVersion?: string,
): void {
this._diag.debug(`Unpatching @pinecone-database/pinecone@${moduleVersion}`);

this._unwrap(moduleExports.Index.prototype, "query");
this._unwrap(moduleExports.Index.prototype, "upsert");
this._unwrap(moduleExports.Index.prototype, "deleteAll");
Expand All @@ -86,6 +93,9 @@ export class PineconeInstrumentation extends InstrumentationBase<any> {
}

private genericWrapper(methodName: string, tracer: Tracer) {
// eslint-disable-next-line @typescript-eslint/no-this-alias
const plugin = this;

// eslint-disable-next-line @typescript-eslint/ban-types
return (original: Function) => {
return function method(this: any, ...args: unknown[]) {
Expand All @@ -98,8 +108,9 @@ export class PineconeInstrumentation extends InstrumentationBase<any> {
return original.apply(this, args);
});
},
// eslint-disable-next-line @typescript-eslint/no-empty-function
() => {},
(e) => {
plugin._diag.error(`Error in Pinecone instrumentation`, e);
},
);
const wrappedPromise = execPromise
.then((result: any) => {
Expand All @@ -125,6 +136,9 @@ export class PineconeInstrumentation extends InstrumentationBase<any> {
}

private queryWrapper(tracer: Tracer) {
// eslint-disable-next-line @typescript-eslint/no-this-alias
const plugin = this;

// eslint-disable-next-line @typescript-eslint/ban-types
return (original: Function) => {
return function method(this: any, ...args: unknown[]) {
Expand Down Expand Up @@ -164,8 +178,11 @@ export class PineconeInstrumentation extends InstrumentationBase<any> {
return original.apply(this, args);
});
},
// eslint-disable-next-line @typescript-eslint/no-empty-function
() => {},
(e) => {
if (e) {
plugin._diag.error(`Error in Pinecone instrumentation`, e);
}
},
);
const wrappedPromise = execPromise
.then((result: any) => {
Expand Down
Loading
Loading