11'use strict'
2-
32const { addHook } = require ( './helpers/instrument' )
3+
44const shimmer = require ( '../../datadog-shimmer' )
55
66const tracingChannel = require ( 'dc-polyfill' ) . tracingChannel
77
8- const invokeTracingChannel = tracingChannel ( 'apm:langchain:invoke' )
9-
10- function wrapLangChainPromise ( fn , type , namespace = [ ] ) {
11- return function ( ) {
12- if ( ! invokeTracingChannel . start . hasSubscribers ) {
13- return fn . apply ( this , arguments )
14- }
15-
16- // Runnable interfaces have an `lc_namespace` property
17- const ns = this . lc_namespace || namespace
18- const resource = [ ...ns , this . constructor . name ] . join ( '.' )
19-
20- const ctx = {
21- args : arguments ,
22- instance : this ,
23- type,
24- resource
8+ function wrap ( obj , name , channelName , namespace ) {
9+ const channel = tracingChannel ( channelName )
10+ shimmer . wrap ( obj , name , function ( original ) {
11+ return function ( ) {
12+ if ( ! channel . start . hasSubscribers ) {
13+ return original . apply ( this , arguments )
14+ }
15+ const ctx = { self : this , arguments }
16+ if ( namespace ) {
17+ ctx . namespace = namespace
18+ }
19+ return channel . tracePromise ( original , ctx , this , ...arguments )
2520 }
26-
27- return invokeTracingChannel . tracePromise ( fn , ctx , this , ...arguments )
28- }
21+ } )
2922}
3023
3124// langchain compiles into ESM and CommonJS, with ESM being the default and landing in the `.js` files
@@ -35,9 +28,10 @@ const extensions = ['js', 'cjs']
3528
3629for ( const extension of extensions ) {
3730 addHook ( { name : '@langchain/core' , file : `dist/runnables/base.${ extension } , versions : [ '>=0.1' ] } , exports => {
38- const RunnableSequence = exports . RunnableSequence
39- shimmer . wrap ( RunnableSequence . prototype , 'invoke' , invoke => wrapLangChainPromise ( invoke , 'chain' ) )
40- shimmer . wrap ( RunnableSequence . prototype , 'batch' , batch => wrapLangChainPromise ( batch , 'chain' ) )
31+ if ( extension === 'cjs' ) {
32+ wrap ( exports . RunnableSequence . prototype , 'invoke' , 'orchestrion:@langchain/core:RunnableSequence_invoke' )
33+ wrap ( exports . RunnableSequence . prototype , 'batch' , 'orchestrion:@langchain/core:RunnableSequence_batch' )
34+ }
4135 return exports
4236 } )
4337
@@ -46,51 +40,53 @@ for (const extension of extensions) {
4640 file : `dist/language_models/chat_models.${ extension } ,
4741 versions : [ '>=0.1' ]
4842 } , exports => {
49- const BaseChatModel = exports . BaseChatModel
50- shimmer . wrap (
51- BaseChatModel . prototype ,
52- 'generate' ,
53- generate => wrapLangChainPromise ( generate , 'chat_model' )
54- )
43+ if ( extension === 'cjs' ) {
44+ wrap ( exports . BaseChatModel . prototype , 'generate' , 'orchestrion:@langchain/core:BaseChatModel_generate' )
45+ }
5546 return exports
5647 } )
5748
5849 addHook ( { name : '@langchain/core' , file : `dist/language_models/llms.${ extension } , versions : [ '>=0.1' ] } , exports => {
59- const BaseLLM = exports . BaseLLM
60- shimmer . wrap ( BaseLLM . prototype , 'generate' , generate => wrapLangChainPromise ( generate , 'llm' ) )
50+ if ( extension === 'cjs' ) {
51+ wrap ( exports . BaseLLM . prototype , 'generate' , 'orchestrion:@langchain/core:BaseLLM_generate' )
52+ }
6153 return exports
6254 } )
6355
6456 addHook ( { name : '@langchain/core' , file : `dist/embeddings.${ extension } , versions : [ '>=0.1' ] } , exports => {
65- // we cannot patch the prototype of the Embeddings class directly
66- // this is because the "abstract class Embeddings" is transpiled from TypeScript to not include abstract functions
67- // thus, we patch the exported class directly instead instead.
57+ if ( extension === 'cjs' ) {
58+ shimmer . wrap ( exports , 'Embeddings' , Embeddings => {
59+ return class extends Embeddings {
60+ constructor ( ...args ) {
61+ super ( ...args )
62+
63+ const namespace = [ 'langchain' , 'embeddings' ]
6864
69- shimmer . wrap ( exports , 'Embeddings' , Embeddings => {
70- return class extends Embeddings {
71- constructor ( ...args ) {
72- super ( ...args )
65+ if ( this . constructor . name === 'OpenAIEmbeddings' ) {
66+ namespace . push ( 'openai' )
67+ }
7368
69+ wrap ( this , 'embedQuery' , 'apm:@langchain/core:Embeddings_embedQuery' , namespace )
70+ wrap ( this , 'embedDocuments' , 'apm:@langchain/core:Embeddings_embedDocuments' , namespace )
71+ }
72+ }
73+ } )
74+ } else {
75+ const channel = tracingChannel ( 'orchestrion:@langchain/core:Embeddings_constructor' )
76+ channel . subscribe ( {
77+ end ( ctx ) {
78+ const { self } = ctx
7479 const namespace = [ 'langchain' , 'embeddings' ]
7580
76- // when originally implemented, we only wrapped OpenAI embeddings
77- // these embeddings had the resource name of `langchain.embeddings.openai.OpenAIEmbeddings`
78- // we need to make sure `openai` is appended to the resource name until a new tracer major version
79- if ( this . constructor . name === 'OpenAIEmbeddings' ) {
81+ if ( self . constructor . name === 'OpenAIEmbeddings' ) {
8082 namespace . push ( 'openai' )
8183 }
8284
83- shimmer . wrap ( this , 'embedQuery' , embedQuery => wrapLangChainPromise ( embedQuery , 'embedding' , namespace ) )
84- shimmer . wrap ( this , 'embedDocuments' ,
85- embedDocuments => wrapLangChainPromise ( embedDocuments , 'embedding' , namespace ) )
85+ wrap ( self , 'embedQuery' , 'apm:@langchain/core:Embeddings_embedQuery' , namespace )
86+ wrap ( self , 'embedDocuments' , 'apm:@langchain/core:Embeddings_embedDocuments' , namespace )
8687 }
87-
88- static [ Symbol . hasInstance ] ( instance ) {
89- return instance instanceof Embeddings
90- }
91- }
92- } )
93-
88+ } )
89+ }
9490 return exports
9591 } )
9692}
0 commit comments