Skip to content

Commit 7950cb9

Browse files
nirganirga
authored
fix: do not patch modules (#146)
1 parent 35b5fa8 commit 7950cb9

File tree

9 files changed

+33
-189
lines changed

9 files changed

+33
-189
lines changed

‎packages/instrumentation-azure/src/instrumentation.ts

+3-21
Original file line numberDiff line numberDiff line change
@@ -50,13 +50,7 @@ export class AzureOpenAIInstrumentation extends InstrumentationBase<any> {
5050
super.setConfig(config);
5151
}
5252

53-
public manuallyInstrument(
54-
module: typeof azure & { openLLMetryPatched?: boolean },
55-
) {
56-
if (module.openLLMetryPatched) {
57-
return;
58-
}
59-
53+
public manuallyInstrument(module: typeof azure) {
6054
this._wrap(
6155
module.OpenAIClient.prototype,
6256
"getChatCompletions",
@@ -80,15 +74,7 @@ export class AzureOpenAIInstrumentation extends InstrumentationBase<any> {
8074
return module;
8175
}
8276

83-
private patch(
84-
moduleExports: typeof azure & { openLLMetryPatched?: boolean },
85-
) {
86-
if (moduleExports.openLLMetryPatched) {
87-
return moduleExports;
88-
}
89-
90-
moduleExports.openLLMetryPatched = true;
91-
77+
private patch(moduleExports: typeof azure) {
9278
this._wrap(
9379
moduleExports.OpenAIClient.prototype,
9480
"getChatCompletions",
@@ -102,11 +88,7 @@ export class AzureOpenAIInstrumentation extends InstrumentationBase<any> {
10288
return moduleExports;
10389
}
10490

105-
private unpatch(
106-
moduleExports: typeof azure & { openLLMetryPatched?: boolean },
107-
): void {
108-
moduleExports.openLLMetryPatched = false;
109-
91+
private unpatch(moduleExports: typeof azure): void {
11092
this._unwrap(moduleExports.OpenAIClient.prototype, "getChatCompletions");
11193
this._unwrap(moduleExports.OpenAIClient.prototype, "getCompletions");
11294
}

‎packages/instrumentation-bedrock/src/instrumentation.ts

+3-19
Original file line numberDiff line numberDiff line change
@@ -57,29 +57,15 @@ export class BedrockInstrumentation extends InstrumentationBase<any> {
5757
return module;
5858
}
5959

60-
public manuallyInstrument(
61-
module: typeof bedrock & { openLLMetryPatched?: boolean },
62-
) {
63-
if (module.openLLMetryPatched) {
64-
return;
65-
}
66-
67-
module.openLLMetryPatched = true;
68-
60+
public manuallyInstrument(module: typeof bedrock) {
6961
this._wrap(
7062
module.BedrockRuntimeClient.prototype,
7163
"send",
7264
this.wrapperMethod(),
7365
);
7466
}
7567

76-
private wrap(module: typeof bedrock & { openLLMetryPatched?: boolean }) {
77-
if (module.openLLMetryPatched) {
78-
return module;
79-
}
80-
81-
module.openLLMetryPatched = true;
82-
68+
private wrap(module: typeof bedrock) {
8369
this._wrap(
8470
module.BedrockRuntimeClient.prototype,
8571
"send",
@@ -89,9 +75,7 @@ export class BedrockInstrumentation extends InstrumentationBase<any> {
8975
return module;
9076
}
9177

92-
private unwrap(module: typeof bedrock & { openLLMetryPatched?: boolean }) {
93-
module.openLLMetryPatched = false;
94-
78+
private unwrap(module: typeof bedrock) {
9579
this._unwrap(module.BedrockRuntimeClient.prototype, "send");
9680
}
9781

‎packages/instrumentation-cohere/src/instrumentation.ts

+3-13
Original file line numberDiff line numberDiff line change
@@ -59,19 +59,11 @@ export class CohereInstrumentation extends InstrumentationBase<any> {
5959
return module;
6060
}
6161

62-
public manuallyInstrument(
63-
module: typeof cohere & { openLLMetryPatched?: boolean },
64-
) {
62+
public manuallyInstrument(module: typeof cohere) {
6563
this.wrap(module);
6664
}
6765

68-
private wrap(module: typeof cohere & { openLLMetryPatched?: boolean }) {
69-
if (module.openLLMetryPatched) {
70-
return module;
71-
}
72-
73-
module.openLLMetryPatched = true;
74-
66+
private wrap(module: typeof cohere) {
7567
this._wrap(
7668
module.CohereClient.prototype,
7769
"generate",
@@ -101,9 +93,7 @@ export class CohereInstrumentation extends InstrumentationBase<any> {
10193
return module;
10294
}
10395

104-
private unwrap(module: typeof cohere & { openLLMetryPatched?: boolean }) {
105-
module.openLLMetryPatched = false;
106-
96+
private unwrap(module: typeof cohere) {
10797
this._unwrap(module.CohereClient.prototype, "generateStream");
10898
this._unwrap(module.CohereClient.prototype, "chat");
10999
this._unwrap(module.CohereClient.prototype, "chatStream");

‎packages/instrumentation-langchain/src/instrumentation.ts

+9-39
Original file line numberDiff line numberDiff line change
@@ -37,9 +37,9 @@ export class LangChainInstrumentation extends InstrumentationBase<any> {
3737
agentsModule,
3838
toolsModule,
3939
}: {
40-
chainsModule?: any & { openLLMetryPatched?: boolean };
41-
agentsModule?: any & { openLLMetryPatched?: boolean };
42-
toolsModule?: any & { openLLMetryPatched?: boolean };
40+
chainsModule?: any;
41+
agentsModule?: any;
42+
toolsModule?: any;
4343
}) {
4444
if (chainsModule) {
4545
this.patchChainModule(chainsModule);
@@ -74,13 +74,7 @@ export class LangChainInstrumentation extends InstrumentationBase<any> {
7474
return [chainModule, agentModule, toolsModule];
7575
}
7676

77-
private patchChainModule(
78-
moduleExports: typeof ChainsModule & { openLLMetryPatched?: boolean },
79-
) {
80-
if (moduleExports.openLLMetryPatched) {
81-
return moduleExports;
82-
}
83-
77+
private patchChainModule(moduleExports: typeof ChainsModule) {
8478
this._wrap(
8579
moduleExports.RetrievalQAChain.prototype,
8680
"_call",
@@ -98,13 +92,7 @@ export class LangChainInstrumentation extends InstrumentationBase<any> {
9892
return moduleExports;
9993
}
10094

101-
private patchAgentModule(
102-
moduleExports: typeof AgentsModule & { openLLMetryPatched?: boolean },
103-
) {
104-
if (moduleExports.openLLMetryPatched) {
105-
return moduleExports;
106-
}
107-
95+
private patchAgentModule(moduleExports: typeof AgentsModule) {
10896
this._wrap(
10997
moduleExports.AgentExecutor.prototype,
11098
"_call",
@@ -117,13 +105,7 @@ export class LangChainInstrumentation extends InstrumentationBase<any> {
117105
return moduleExports;
118106
}
119107

120-
private patchToolsModule(
121-
moduleExports: typeof ToolsModule & { openLLMetryPatched?: boolean },
122-
) {
123-
if (moduleExports.openLLMetryPatched) {
124-
return moduleExports;
125-
}
126-
108+
private patchToolsModule(moduleExports: typeof ToolsModule) {
127109
this._wrap(
128110
moduleExports.Tool.prototype,
129111
"call",
@@ -132,30 +114,18 @@ export class LangChainInstrumentation extends InstrumentationBase<any> {
132114
return moduleExports;
133115
}
134116

135-
private unpatchChainModule(
136-
moduleExports: any & { openLLMetryPatched?: boolean },
137-
) {
138-
moduleExports.openLLMetryPatched = false;
139-
117+
private unpatchChainModule(moduleExports: any) {
140118
this._unwrap(moduleExports.RetrievalQAChain.prototype, "_call");
141119
this._unwrap(moduleExports.BaseChain.prototype, "call");
142120
return moduleExports;
143121
}
144122

145-
private unpatchAgentModule(
146-
moduleExports: any & { openLLMetryPatched?: boolean },
147-
) {
148-
moduleExports.openLLMetryPatched = false;
149-
123+
private unpatchAgentModule(moduleExports: any) {
150124
this._unwrap(moduleExports.AgentExecutor.prototype, "_call");
151125
return moduleExports;
152126
}
153127

154-
private unpatchToolsModule(
155-
moduleExports: any & { openLLMetryPatched?: boolean },
156-
) {
157-
moduleExports.openLLMetryPatched = false;
158-
128+
private unpatchToolsModule(moduleExports: any) {
159129
this._unwrap(moduleExports.AgentExecutor.prototype, "_call");
160130
return moduleExports;
161131
}

‎packages/instrumentation-llamaindex/src/instrumentation.ts

+3-17
Original file line numberDiff line numberDiff line change
@@ -39,9 +39,7 @@ export class LlamaIndexInstrumentation extends InstrumentationBase<any> {
3939
super.setConfig(config);
4040
}
4141

42-
public manuallyInstrument(
43-
module: typeof llamaindex & { openLLMetryPatched?: boolean },
44-
) {
42+
public manuallyInstrument(module: typeof llamaindex) {
4543
this.patch(module);
4644
}
4745

@@ -80,15 +78,7 @@ export class LlamaIndexInstrumentation extends InstrumentationBase<any> {
8078
return retriever && (retriever as BaseRetriever).retrieve !== undefined;
8179
}
8280

83-
private patch(
84-
moduleExports: typeof llamaindex & { openLLMetryPatched?: boolean },
85-
) {
86-
if (moduleExports.openLLMetryPatched) {
87-
return moduleExports;
88-
}
89-
90-
moduleExports.openLLMetryPatched = true;
91-
81+
private patch(moduleExports: typeof llamaindex) {
9282
const customLLMInstrumentation = new CustomLLMInstrumentation(
9383
this._config,
9484
() => this.tracer, // this is on purpose. Tracer may change
@@ -168,11 +158,7 @@ export class LlamaIndexInstrumentation extends InstrumentationBase<any> {
168158
return moduleExports;
169159
}
170160

171-
private unpatch(
172-
moduleExports: typeof llamaindex & { openLLMetryPatched?: boolean },
173-
) {
174-
moduleExports.openLLMetryPatched = false;
175-
161+
private unpatch(moduleExports: typeof llamaindex) {
176162
this._unwrap(moduleExports.RetrieverQueryEngine.prototype, "query");
177163

178164
for (const key in moduleExports) {

‎packages/instrumentation-openai/src/instrumentation.ts

+3-21
Original file line numberDiff line numberDiff line change
@@ -55,13 +55,7 @@ export class OpenAIInstrumentation extends InstrumentationBase<any> {
5555
super.setConfig(config);
5656
}
5757

58-
public manuallyInstrument(
59-
module: typeof openai.OpenAI & { openLLMetryPatched?: boolean },
60-
) {
61-
if (module.openLLMetryPatched) {
62-
return;
63-
}
64-
58+
public manuallyInstrument(module: typeof openai.OpenAI) {
6559
// Old version of OpenAI API (v3.1.0)
6660
if ((module as any).OpenAIApi) {
6761
this._wrap(
@@ -98,15 +92,7 @@ export class OpenAIInstrumentation extends InstrumentationBase<any> {
9892
return module;
9993
}
10094

101-
private patch(
102-
moduleExports: typeof openai & { openLLMetryPatched?: boolean },
103-
) {
104-
if (moduleExports.openLLMetryPatched) {
105-
return moduleExports;
106-
}
107-
108-
moduleExports.openLLMetryPatched = true;
109-
95+
private patch(moduleExports: typeof openai) {
11096
// Old version of OpenAI API (v3.1.0)
11197
if ((moduleExports as any).OpenAIApi) {
11298
this._wrap(
@@ -134,11 +120,7 @@ export class OpenAIInstrumentation extends InstrumentationBase<any> {
134120
return moduleExports;
135121
}
136122

137-
private unpatch(
138-
moduleExports: typeof openai & { openLLMetryPatched?: boolean },
139-
): void {
140-
moduleExports.openLLMetryPatched = false;
141-
123+
private unpatch(moduleExports: typeof openai): void {
142124
// Old version of OpenAI API (v3.1.0)
143125
if ((moduleExports as any).OpenAIApi) {
144126
this._unwrap(

‎packages/instrumentation-pinecone/src/instrumentation.ts

+3-17
Original file line numberDiff line numberDiff line change
@@ -33,9 +33,7 @@ export class PineconeInstrumentation extends InstrumentationBase<any> {
3333
super("@traceloop/instrumentation-pinecone", "0.3.0", config);
3434
}
3535

36-
public manuallyInstrument(
37-
module: typeof pinecone & { openLLMetryPatched?: boolean },
38-
) {
36+
public manuallyInstrument(module: typeof pinecone) {
3937
this.patch(module);
4038
}
4139

@@ -49,15 +47,7 @@ export class PineconeInstrumentation extends InstrumentationBase<any> {
4947
return module;
5048
}
5149

52-
private patch(
53-
moduleExports: typeof pinecone & { openLLMetryPatched?: boolean },
54-
) {
55-
if (moduleExports.openLLMetryPatched) {
56-
return moduleExports;
57-
}
58-
59-
moduleExports.openLLMetryPatched = true;
60-
50+
private patch(moduleExports: typeof pinecone) {
6151
this._wrap(
6252
moduleExports.Index.prototype,
6353
"query",
@@ -87,11 +77,7 @@ export class PineconeInstrumentation extends InstrumentationBase<any> {
8777
return moduleExports;
8878
}
8979

90-
private unpatch(
91-
moduleExports: typeof pinecone & { openLLMetryPatched?: boolean },
92-
): void {
93-
moduleExports.openLLMetryPatched = false;
94-
80+
private unpatch(moduleExports: typeof pinecone): void {
9581
this._unwrap(moduleExports.Index.prototype, "query");
9682
this._unwrap(moduleExports.Index.prototype, "upsert");
9783
this._unwrap(moduleExports.Index.prototype, "deleteAll");

‎packages/instrumentation-vertexai/src/aiplatform-instrumentation.ts

+3-21
Original file line numberDiff line numberDiff line change
@@ -57,29 +57,15 @@ export class AIPlatformInstrumentation extends InstrumentationBase<any> {
5757
return aiPlatformModule;
5858
}
5959

60-
public manuallyInstrument(
61-
module: typeof aiplatform & { openLLMetryPatched?: boolean },
62-
) {
63-
if (module.openLLMetryPatched) {
64-
return;
65-
}
66-
67-
module.openLLMetryPatched = true;
68-
60+
public manuallyInstrument(module: typeof aiplatform) {
6961
this._wrap(
7062
module.PredictionServiceClient.prototype,
7163
"predict",
7264
this.wrapperMethod(),
7365
);
7466
}
7567

76-
private wrap(module: typeof aiplatform & { openLLMetryPatched?: boolean }) {
77-
if (module.openLLMetryPatched) {
78-
return module;
79-
}
80-
81-
module.openLLMetryPatched = true;
82-
68+
private wrap(module: typeof aiplatform) {
8369
this._wrap(
8470
module.PredictionServiceClient.prototype,
8571
"predict",
@@ -89,11 +75,7 @@ export class AIPlatformInstrumentation extends InstrumentationBase<any> {
8975
return module;
9076
}
9177

92-
private unwrap(
93-
module: typeof aiplatform & { openLLMetryPatched?: boolean },
94-
): void {
95-
module.openLLMetryPatched = false;
96-
78+
private unwrap(module: typeof aiplatform): void {
9779
this._unwrap(module.PredictionServiceClient.prototype, "predict");
9880
}
9981

0 commit comments

Comments
 (0)