vercel.js 3.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596
  1. import { traceable } from "../traceable.js";
  2. import { _wrapClient } from "./generic.js";
  3. /**
  4. * Wrap a Vercel AI SDK model, enabling automatic LangSmith tracing.
  5. * After wrapping a model, you can use it with the Vercel AI SDK Core
  6. * methods as normal.
  7. *
  8. * @example
  9. * ```ts
  10. * import { anthropic } from "@ai-sdk/anthropic";
  11. * import { streamText } from "ai";
  12. * import { wrapAISDKModel } from "langsmith/wrappers/vercel";
  13. *
  14. * const anthropicModel = anthropic("claude-3-haiku-20240307");
  15. *
  16. * const modelWithTracing = wrapAISDKModel(anthropicModel);
  17. *
  18. * const { textStream } = await streamText({
  19. * model: modelWithTracing,
  20. * prompt: "Write a vegetarian lasagna recipe for 4 people.",
  21. * });
  22. *
  23. * for await (const chunk of textStream) {
  24. * console.log(chunk);
  25. * }
  26. * ```
  27. * @param model An AI SDK model instance.
  28. * @param options LangSmith options.
  29. * @returns
  30. */
  31. export const wrapAISDKModel = (model, options) => {
  32. if (!("doStream" in model) ||
  33. typeof model.doStream !== "function" ||
  34. !("doGenerate" in model) ||
  35. typeof model.doGenerate !== "function") {
  36. throw new Error(`Received invalid input. This version of wrapAISDKModel only supports Vercel LanguageModelV1 instances.`);
  37. }
  38. const runName = options?.name ?? model.constructor?.name;
  39. return new Proxy(model, {
  40. get(target, propKey, receiver) {
  41. const originalValue = target[propKey];
  42. if (typeof originalValue === "function") {
  43. let __finalTracedIteratorKey;
  44. let aggregator;
  45. if (propKey === "doStream") {
  46. __finalTracedIteratorKey = "stream";
  47. aggregator = (chunks) => {
  48. return chunks.reduce((aggregated, chunk) => {
  49. if (chunk.type === "text-delta") {
  50. return {
  51. ...aggregated,
  52. text: aggregated.text + chunk.textDelta,
  53. };
  54. }
  55. else if (chunk.type === "tool-call") {
  56. return {
  57. ...aggregated,
  58. ...chunk,
  59. };
  60. }
  61. else if (chunk.type === "finish") {
  62. return {
  63. ...aggregated,
  64. usage: chunk.usage,
  65. finishReason: chunk.finishReason,
  66. };
  67. }
  68. else {
  69. return aggregated;
  70. }
  71. }, {
  72. text: "",
  73. });
  74. };
  75. }
  76. return traceable(originalValue.bind(target), {
  77. run_type: "llm",
  78. name: runName,
  79. ...options,
  80. __finalTracedIteratorKey,
  81. aggregator,
  82. });
  83. }
  84. else if (originalValue != null &&
  85. !Array.isArray(originalValue) &&
  86. // eslint-disable-next-line no-instanceof/no-instanceof
  87. !(originalValue instanceof Date) &&
  88. typeof originalValue === "object") {
  89. return _wrapClient(originalValue, [runName, propKey.toString()].join("."), options);
  90. }
  91. else {
  92. return Reflect.get(target, propKey, receiver);
  93. }
  94. },
  95. });
  96. };