traceable.cjs 24 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553
  1. "use strict";
  2. Object.defineProperty(exports, "__esModule", { value: true });
  3. exports.ROOT = exports.withRunTree = exports.isTraceableFunction = exports.getCurrentRunTree = void 0;
  4. exports.traceable = traceable;
  5. const node_async_hooks_1 = require("node:async_hooks");
  6. const run_trees_js_1 = require("./run_trees.cjs");
  7. const env_js_1 = require("./env.cjs");
  8. const traceable_js_1 = require("./singletons/traceable.cjs");
  9. const constants_js_1 = require("./singletons/constants.cjs");
  10. const asserts_js_1 = require("./utils/asserts.cjs");
  11. traceable_js_1.AsyncLocalStorageProviderSingleton.initializeGlobalInstance(new node_async_hooks_1.AsyncLocalStorage());
  12. const runInputsToMap = (rawInputs) => {
  13. const firstInput = rawInputs[0];
  14. let inputs;
  15. if (firstInput == null) {
  16. inputs = {};
  17. }
  18. else if (rawInputs.length > 1) {
  19. inputs = { args: rawInputs };
  20. }
  21. else if ((0, asserts_js_1.isKVMap)(firstInput)) {
  22. inputs = firstInput;
  23. }
  24. else {
  25. inputs = { input: firstInput };
  26. }
  27. return inputs;
  28. };
  29. const handleRunInputs = (inputs, processInputs) => {
  30. try {
  31. return processInputs(inputs);
  32. }
  33. catch (e) {
  34. console.error("Error occurred during processInputs. Sending raw inputs:", e);
  35. return inputs;
  36. }
  37. };
  38. const handleRunOutputs = (rawOutputs, processOutputs) => {
  39. let outputs;
  40. if ((0, asserts_js_1.isKVMap)(rawOutputs)) {
  41. outputs = rawOutputs;
  42. }
  43. else {
  44. outputs = { outputs: rawOutputs };
  45. }
  46. try {
  47. return processOutputs(outputs);
  48. }
  49. catch (e) {
  50. console.error("Error occurred during processOutputs. Sending raw outputs:", e);
  51. return outputs;
  52. }
  53. };
  54. const handleRunAttachments = (rawInputs, extractAttachments) => {
  55. if (!extractAttachments) {
  56. return [undefined, rawInputs];
  57. }
  58. try {
  59. const [attachments, remainingArgs] = extractAttachments(...rawInputs);
  60. return [attachments, remainingArgs];
  61. }
  62. catch (e) {
  63. console.error("Error occurred during extractAttachments:", e);
  64. return [undefined, rawInputs];
  65. }
  66. };
  67. const getTracingRunTree = (runTree, inputs, getInvocationParams, processInputs, extractAttachments) => {
  68. if (!(0, env_js_1.isTracingEnabled)(runTree.tracingEnabled)) {
  69. return undefined;
  70. }
  71. const [attached, args] = handleRunAttachments(inputs, extractAttachments);
  72. runTree.attachments = attached;
  73. runTree.inputs = handleRunInputs(args, processInputs);
  74. const invocationParams = getInvocationParams?.(...inputs);
  75. if (invocationParams != null) {
  76. runTree.extra ??= {};
  77. runTree.extra.metadata = {
  78. ...invocationParams,
  79. ...runTree.extra.metadata,
  80. };
  81. }
  82. return runTree;
  83. };
  84. // idea: store the state of the promise outside
  85. // but only when the promise is "consumed"
  86. const getSerializablePromise = (arg) => {
  87. const proxyState = { current: undefined };
  88. const promiseProxy = new Proxy(arg, {
  89. get(target, prop, receiver) {
  90. if (prop === "then") {
  91. const boundThen = arg[prop].bind(arg);
  92. return (resolve, reject = (x) => {
  93. throw x;
  94. }) => {
  95. return boundThen((value) => {
  96. proxyState.current = ["resolve", value];
  97. return resolve(value);
  98. }, (error) => {
  99. proxyState.current = ["reject", error];
  100. return reject(error);
  101. });
  102. };
  103. }
  104. if (prop === "catch") {
  105. const boundCatch = arg[prop].bind(arg);
  106. return (reject) => {
  107. return boundCatch((error) => {
  108. proxyState.current = ["reject", error];
  109. return reject(error);
  110. });
  111. };
  112. }
  113. if (prop === "toJSON") {
  114. return () => {
  115. if (!proxyState.current)
  116. return undefined;
  117. const [type, value] = proxyState.current ?? [];
  118. if (type === "resolve")
  119. return value;
  120. return { error: value };
  121. };
  122. }
  123. return Reflect.get(target, prop, receiver);
  124. },
  125. });
  126. return promiseProxy;
  127. };
  128. const convertSerializableArg = (arg) => {
  129. if ((0, asserts_js_1.isReadableStream)(arg)) {
  130. const proxyState = [];
  131. const transform = new TransformStream({
  132. start: () => void 0,
  133. transform: (chunk, controller) => {
  134. proxyState.push(chunk);
  135. controller.enqueue(chunk);
  136. },
  137. flush: () => void 0,
  138. });
  139. const pipeThrough = arg.pipeThrough(transform);
  140. Object.assign(pipeThrough, { toJSON: () => proxyState });
  141. return pipeThrough;
  142. }
  143. if ((0, asserts_js_1.isAsyncIterable)(arg)) {
  144. const proxyState = { current: [] };
  145. return new Proxy(arg, {
  146. get(target, prop, receiver) {
  147. if (prop === Symbol.asyncIterator) {
  148. return () => {
  149. const boundIterator = arg[Symbol.asyncIterator].bind(arg);
  150. const iterator = boundIterator();
  151. return new Proxy(iterator, {
  152. get(target, prop, receiver) {
  153. if (prop === "next" || prop === "return" || prop === "throw") {
  154. const bound = iterator.next.bind(iterator);
  155. return (...args) => {
  156. // @ts-expect-error TS cannot infer the argument types for the bound function
  157. const wrapped = getSerializablePromise(bound(...args));
  158. proxyState.current.push(wrapped);
  159. return wrapped;
  160. };
  161. }
  162. if (prop === "return" || prop === "throw") {
  163. return iterator.next.bind(iterator);
  164. }
  165. return Reflect.get(target, prop, receiver);
  166. },
  167. });
  168. };
  169. }
  170. if (prop === "toJSON") {
  171. return () => {
  172. const onlyNexts = proxyState.current;
  173. const serialized = onlyNexts.map((next) => next.toJSON());
  174. const chunks = serialized.reduce((memo, next) => {
  175. if (next?.value)
  176. memo.push(next.value);
  177. return memo;
  178. }, []);
  179. return chunks;
  180. };
  181. }
  182. return Reflect.get(target, prop, receiver);
  183. },
  184. });
  185. }
  186. if (!Array.isArray(arg) && (0, asserts_js_1.isIteratorLike)(arg)) {
  187. const proxyState = [];
  188. return new Proxy(arg, {
  189. get(target, prop, receiver) {
  190. if (prop === "next" || prop === "return" || prop === "throw") {
  191. const bound = arg[prop]?.bind(arg);
  192. return (...args) => {
  193. const next = bound?.(...args);
  194. if (next != null)
  195. proxyState.push(next);
  196. return next;
  197. };
  198. }
  199. if (prop === "toJSON") {
  200. return () => {
  201. const chunks = proxyState.reduce((memo, next) => {
  202. if (next.value)
  203. memo.push(next.value);
  204. return memo;
  205. }, []);
  206. return chunks;
  207. };
  208. }
  209. return Reflect.get(target, prop, receiver);
  210. },
  211. });
  212. }
  213. if ((0, asserts_js_1.isThenable)(arg)) {
  214. return getSerializablePromise(arg);
  215. }
  216. return arg;
  217. };
  218. /**
  219. * Higher-order function that takes function as input and returns a
  220. * "TraceableFunction" - a wrapped version of the input that
  221. * automatically handles tracing. If the returned traceable function calls any
  222. * traceable functions, those are automatically traced as well.
  223. *
  224. * The returned TraceableFunction can accept a run tree or run tree config as
  225. * its first argument. If omitted, it will default to the caller's run tree,
  226. * or will be treated as a root run.
  227. *
  228. * @param wrappedFunc Targeted function to be traced
  229. * @param config Additional metadata such as name, tags or providing
  230. * a custom LangSmith client instance
  231. */
  232. // eslint-disable-next-line @typescript-eslint/no-explicit-any
  233. function traceable(wrappedFunc, config) {
  234. const { aggregator, argsConfigPath, __finalTracedIteratorKey, processInputs, processOutputs, extractAttachments, ...runTreeConfig } = config ?? {};
  235. const processInputsFn = processInputs ?? ((x) => x);
  236. const processOutputsFn = processOutputs ?? ((x) => x);
  237. const extractAttachmentsFn = extractAttachments ?? ((...x) => [undefined, runInputsToMap(x)]);
  238. const traceableFunc = (...args) => {
  239. let ensuredConfig;
  240. try {
  241. let runtimeConfig;
  242. if (argsConfigPath) {
  243. const [index, path] = argsConfigPath;
  244. if (index === args.length - 1 && !path) {
  245. runtimeConfig = args.pop();
  246. }
  247. else if (index <= args.length &&
  248. typeof args[index] === "object" &&
  249. args[index] !== null) {
  250. if (path) {
  251. const { [path]: extracted, ...rest } = args[index];
  252. runtimeConfig = extracted;
  253. args[index] = rest;
  254. }
  255. else {
  256. runtimeConfig = args[index];
  257. args.splice(index, 1);
  258. }
  259. }
  260. }
  261. ensuredConfig = {
  262. name: wrappedFunc.name || "<lambda>",
  263. ...runTreeConfig,
  264. ...runtimeConfig,
  265. tags: [
  266. ...new Set([
  267. ...(runTreeConfig?.tags ?? []),
  268. ...(runtimeConfig?.tags ?? []),
  269. ]),
  270. ],
  271. metadata: {
  272. ...runTreeConfig?.metadata,
  273. ...runtimeConfig?.metadata,
  274. },
  275. };
  276. }
  277. catch (err) {
  278. console.warn(`Failed to extract runtime config from args for ${runTreeConfig?.name ?? wrappedFunc.name}`, err);
  279. ensuredConfig = {
  280. name: wrappedFunc.name || "<lambda>",
  281. ...runTreeConfig,
  282. };
  283. }
  284. const asyncLocalStorage = traceable_js_1.AsyncLocalStorageProviderSingleton.getInstance();
  285. // TODO: deal with possible nested promises and async iterables
  286. const processedArgs = args;
  287. for (let i = 0; i < processedArgs.length; i++) {
  288. processedArgs[i] = convertSerializableArg(processedArgs[i]);
  289. }
  290. const [currentRunTree, rawInputs] = (() => {
  291. const [firstArg, ...restArgs] = processedArgs;
  292. // used for handoff between LangChain.JS and traceable functions
  293. if ((0, run_trees_js_1.isRunnableConfigLike)(firstArg)) {
  294. return [
  295. getTracingRunTree(run_trees_js_1.RunTree.fromRunnableConfig(firstArg, ensuredConfig), restArgs, config?.getInvocationParams, processInputsFn, extractAttachmentsFn),
  296. restArgs,
  297. ];
  298. }
  299. // deprecated: legacy CallbackManagerRunTree used in runOnDataset
  300. // override ALS and do not pass-through the run tree
  301. if ((0, run_trees_js_1.isRunTree)(firstArg) &&
  302. "callbackManager" in firstArg &&
  303. firstArg.callbackManager != null) {
  304. return [firstArg, restArgs];
  305. }
  306. // when ALS is unreliable, users can manually
  307. // pass in the run tree
  308. if (firstArg === traceable_js_1.ROOT || (0, run_trees_js_1.isRunTree)(firstArg)) {
  309. const currentRunTree = getTracingRunTree(firstArg === traceable_js_1.ROOT
  310. ? new run_trees_js_1.RunTree(ensuredConfig)
  311. : firstArg.createChild(ensuredConfig), restArgs, config?.getInvocationParams, processInputsFn, extractAttachmentsFn);
  312. return [currentRunTree, [currentRunTree, ...restArgs]];
  313. }
  314. // Node.JS uses AsyncLocalStorage (ALS) and AsyncResource
  315. // to allow storing context
  316. const prevRunFromStore = asyncLocalStorage.getStore();
  317. if ((0, run_trees_js_1.isRunTree)(prevRunFromStore)) {
  318. return [
  319. getTracingRunTree(prevRunFromStore.createChild(ensuredConfig), processedArgs, config?.getInvocationParams, processInputsFn, extractAttachmentsFn),
  320. processedArgs,
  321. ];
  322. }
  323. const currentRunTree = getTracingRunTree(new run_trees_js_1.RunTree(ensuredConfig), processedArgs, config?.getInvocationParams, processInputsFn, extractAttachmentsFn);
  324. // If a context var is set by LangChain outside of a traceable,
  325. // it will be an object with a single property and we should copy
  326. // context vars over into the new run tree.
  327. if (prevRunFromStore !== undefined &&
  328. constants_js_1._LC_CONTEXT_VARIABLES_KEY in prevRunFromStore) {
  329. // eslint-disable-next-line @typescript-eslint/no-explicit-any
  330. currentRunTree[constants_js_1._LC_CONTEXT_VARIABLES_KEY] =
  331. prevRunFromStore[constants_js_1._LC_CONTEXT_VARIABLES_KEY];
  332. }
  333. return [currentRunTree, processedArgs];
  334. })();
  335. return asyncLocalStorage.run(currentRunTree, () => {
  336. const postRunPromise = currentRunTree?.postRun();
  337. async function handleChunks(chunks) {
  338. if (aggregator !== undefined) {
  339. try {
  340. return await aggregator(chunks);
  341. }
  342. catch (e) {
  343. console.error(`[ERROR]: LangSmith aggregation failed: `, e);
  344. }
  345. }
  346. return chunks;
  347. }
  348. function tapReadableStreamForTracing(stream, snapshot) {
  349. const reader = stream.getReader();
  350. let finished = false;
  351. const chunks = [];
  352. const tappedStream = new ReadableStream({
  353. async start(controller) {
  354. // eslint-disable-next-line no-constant-condition
  355. while (true) {
  356. const result = await (snapshot
  357. ? snapshot(() => reader.read())
  358. : reader.read());
  359. if (result.done) {
  360. finished = true;
  361. await currentRunTree?.end(handleRunOutputs(await handleChunks(chunks), processOutputsFn));
  362. await handleEnd();
  363. controller.close();
  364. break;
  365. }
  366. chunks.push(result.value);
  367. // Add new_token event for streaming LLM runs
  368. if (currentRunTree && currentRunTree.run_type === "llm") {
  369. currentRunTree.addEvent({
  370. name: "new_token",
  371. kwargs: { token: result.value },
  372. });
  373. }
  374. controller.enqueue(result.value);
  375. }
  376. },
  377. async cancel(reason) {
  378. if (!finished)
  379. await currentRunTree?.end(undefined, "Cancelled");
  380. await currentRunTree?.end(handleRunOutputs(await handleChunks(chunks), processOutputsFn));
  381. await handleEnd();
  382. return reader.cancel(reason);
  383. },
  384. });
  385. return tappedStream;
  386. }
  387. async function* wrapAsyncIteratorForTracing(iterator, snapshot) {
  388. let finished = false;
  389. const chunks = [];
  390. try {
  391. while (true) {
  392. const { value, done } = await (snapshot
  393. ? snapshot(() => iterator.next())
  394. : iterator.next());
  395. if (done) {
  396. finished = true;
  397. break;
  398. }
  399. chunks.push(value);
  400. // Add new_token event for streaming LLM runs
  401. if (currentRunTree && currentRunTree.run_type === "llm") {
  402. currentRunTree.addEvent({
  403. name: "new_token",
  404. kwargs: { token: value },
  405. });
  406. }
  407. yield value;
  408. }
  409. }
  410. catch (e) {
  411. await currentRunTree?.end(undefined, String(e));
  412. throw e;
  413. }
  414. finally {
  415. if (!finished)
  416. await currentRunTree?.end(undefined, "Cancelled");
  417. await currentRunTree?.end(handleRunOutputs(await handleChunks(chunks), processOutputsFn));
  418. await handleEnd();
  419. }
  420. }
  421. function wrapAsyncGeneratorForTracing(iterable, snapshot) {
  422. if ((0, asserts_js_1.isReadableStream)(iterable)) {
  423. return tapReadableStreamForTracing(iterable, snapshot);
  424. }
  425. const iterator = iterable[Symbol.asyncIterator]();
  426. const wrappedIterator = wrapAsyncIteratorForTracing(iterator, snapshot);
  427. iterable[Symbol.asyncIterator] = () => wrappedIterator;
  428. return iterable;
  429. }
  430. async function handleEnd() {
  431. const onEnd = config?.on_end;
  432. if (onEnd) {
  433. if (!currentRunTree) {
  434. console.warn("Can not call 'on_end' if currentRunTree is undefined");
  435. }
  436. else {
  437. onEnd(currentRunTree);
  438. }
  439. }
  440. await postRunPromise;
  441. await currentRunTree?.patchRun();
  442. }
  443. function gatherAll(iterator) {
  444. const chunks = [];
  445. // eslint-disable-next-line no-constant-condition
  446. while (true) {
  447. const next = iterator.next();
  448. chunks.push(next);
  449. if (next.done)
  450. break;
  451. }
  452. return chunks;
  453. }
  454. let returnValue;
  455. try {
  456. returnValue = wrappedFunc(...rawInputs);
  457. }
  458. catch (err) {
  459. returnValue = Promise.reject(err);
  460. }
  461. if ((0, asserts_js_1.isAsyncIterable)(returnValue)) {
  462. const snapshot = node_async_hooks_1.AsyncLocalStorage.snapshot();
  463. return wrapAsyncGeneratorForTracing(returnValue, snapshot);
  464. }
  465. if (!Array.isArray(returnValue) &&
  466. typeof returnValue === "object" &&
  467. returnValue != null &&
  468. __finalTracedIteratorKey !== undefined &&
  469. (0, asserts_js_1.isAsyncIterable)(returnValue[__finalTracedIteratorKey])) {
  470. const snapshot = node_async_hooks_1.AsyncLocalStorage.snapshot();
  471. return {
  472. ...returnValue,
  473. [__finalTracedIteratorKey]: wrapAsyncGeneratorForTracing(returnValue[__finalTracedIteratorKey], snapshot),
  474. };
  475. }
  476. const tracedPromise = new Promise((resolve, reject) => {
  477. Promise.resolve(returnValue)
  478. .then(async (rawOutput) => {
  479. if ((0, asserts_js_1.isAsyncIterable)(rawOutput)) {
  480. const snapshot = node_async_hooks_1.AsyncLocalStorage.snapshot();
  481. return resolve(wrapAsyncGeneratorForTracing(rawOutput, snapshot));
  482. }
  483. if (!Array.isArray(rawOutput) &&
  484. typeof rawOutput === "object" &&
  485. rawOutput != null &&
  486. __finalTracedIteratorKey !== undefined &&
  487. (0, asserts_js_1.isAsyncIterable)(rawOutput[__finalTracedIteratorKey])) {
  488. const snapshot = node_async_hooks_1.AsyncLocalStorage.snapshot();
  489. return {
  490. ...rawOutput,
  491. [__finalTracedIteratorKey]: wrapAsyncGeneratorForTracing(rawOutput[__finalTracedIteratorKey], snapshot),
  492. };
  493. }
  494. if ((0, asserts_js_1.isGenerator)(wrappedFunc) && (0, asserts_js_1.isIteratorLike)(rawOutput)) {
  495. const chunks = gatherAll(rawOutput);
  496. try {
  497. await currentRunTree?.end(handleRunOutputs(await handleChunks(chunks.reduce((memo, { value, done }) => {
  498. if (!done || typeof value !== "undefined") {
  499. memo.push(value);
  500. }
  501. return memo;
  502. }, [])), processOutputsFn));
  503. await handleEnd();
  504. }
  505. catch (e) {
  506. console.error("Error occurred during handleEnd:", e);
  507. }
  508. return (function* () {
  509. for (const ret of chunks) {
  510. if (ret.done)
  511. return ret.value;
  512. yield ret.value;
  513. }
  514. })();
  515. }
  516. try {
  517. await currentRunTree?.end(handleRunOutputs(rawOutput, processOutputsFn));
  518. await handleEnd();
  519. }
  520. finally {
  521. // eslint-disable-next-line no-unsafe-finally
  522. return rawOutput;
  523. }
  524. }, async (error) => {
  525. await currentRunTree?.end(undefined, String(error));
  526. await handleEnd();
  527. throw error;
  528. })
  529. .then(resolve, reject);
  530. });
  531. if (typeof returnValue !== "object" || returnValue === null) {
  532. return tracedPromise;
  533. }
  534. return new Proxy(returnValue, {
  535. get(target, prop, receiver) {
  536. if ((0, asserts_js_1.isPromiseMethod)(prop)) {
  537. return tracedPromise[prop].bind(tracedPromise);
  538. }
  539. return Reflect.get(target, prop, receiver);
  540. },
  541. });
  542. });
  543. };
  544. Object.defineProperty(traceableFunc, "langsmith:traceable", {
  545. value: runTreeConfig,
  546. });
  547. return traceableFunc;
  548. }
  549. var traceable_js_2 = require("./singletons/traceable.cjs");
  550. Object.defineProperty(exports, "getCurrentRunTree", { enumerable: true, get: function () { return traceable_js_2.getCurrentRunTree; } });
  551. Object.defineProperty(exports, "isTraceableFunction", { enumerable: true, get: function () { return traceable_js_2.isTraceableFunction; } });
  552. Object.defineProperty(exports, "withRunTree", { enumerable: true, get: function () { return traceable_js_2.withRunTree; } });
  553. Object.defineProperty(exports, "ROOT", { enumerable: true, get: function () { return traceable_js_2.ROOT; } });