ai 6.0.32 → 6.0.34

package/src/rerank/rerank.ts

@@ -0,0 +1,237 @@
+import { JSONObject, RerankingModelV3CallOptions } from '@ai-sdk/provider';
+import { ProviderOptions } from '@ai-sdk/provider-utils';
+import { prepareRetries } from '../../src/util/prepare-retries';
+import { assembleOperationName } from '../telemetry/assemble-operation-name';
+import { getBaseTelemetryAttributes } from '../telemetry/get-base-telemetry-attributes';
+import { getTracer } from '../telemetry/get-tracer';
+import { recordSpan } from '../telemetry/record-span';
+import { selectTelemetryAttributes } from '../telemetry/select-telemetry-attributes';
+import { TelemetrySettings } from '../telemetry/telemetry-settings';
+import { RerankingModel } from '../types';
+import { RerankResult } from './rerank-result';
+import { logWarnings } from '../logger/log-warnings';
+
+/**
+Rerank documents using an reranking model. The type of the value is defined by the reranking model.
+
+@param model - The Reranking model to use.
+@param documents - The documents that should be reranking.
+@param query - The query is a string that represents the query to rerank the documents against.
+@param topN - Top n documents to rerank.
+
+@param maxRetries - Maximum number of retries. Set to 0 to disable retries. Default: 2.
+@param abortSignal - An optional abort signal that can be used to cancel the call.
+@param headers - Additional HTTP headers to be sent with the request. Only applicable for HTTP-based providers.
+
+@returns A result object that contains the reranked documents, the reranked indices, and additional information.
+ */
+export async function rerank<VALUE extends JSONObject | string>({
+  model,
+  documents,
+  query,
+  topN,
+  maxRetries: maxRetriesArg,
+  abortSignal,
+  headers,
+  providerOptions,
+  experimental_telemetry: telemetry,
+}: {
+  /**
+The reranking model to use.
+  */
+  model: RerankingModel;
+
+  /**
+   * The documents that should be reranked.
+   */
+  documents: Array<VALUE>;
+
+  /**
+The query is a string that represents the query to rerank the documents against.
+   */
+  query: string;
+
+  /**
+   * Number of top documents to return.
+   */
+  topN?: number;
+
+  /**
+Maximum number of retries per reranking model call. Set to 0 to disable retries.
+
+@default 2
+   */
+  maxRetries?: number;
+
+  /**
+Abort signal.
+ */
+  abortSignal?: AbortSignal;
+
+  /**
+Additional headers to include in the request.
+Only applicable for HTTP-based providers.
+ */
+  headers?: Record<string, string>;
+
+  /**
+   * Optional telemetry configuration (experimental).
+   */
+  experimental_telemetry?: TelemetrySettings;
+
+  /**
+    Additional provider-specific options. They are passed through
+    to the provider from the AI SDK and enable provider-specific
+    functionality that can be fully encapsulated in the provider.
+    */
+  providerOptions?: ProviderOptions;
+}): Promise<RerankResult<VALUE>> {
+  if (documents.length === 0) {
+    return new DefaultRerankResult({
+      originalDocuments: [],
+      ranking: [],
+      providerMetadata: undefined,
+      response: {
+        timestamp: new Date(),
+        modelId: model.modelId,
+      },
+    });
+  }
+
+  const { maxRetries, retry } = prepareRetries({
+    maxRetries: maxRetriesArg,
+    abortSignal,
+  });
+
+  // detect the type of the documents:
+  const documentsToSend: RerankingModelV3CallOptions['documents'] =
+    typeof documents[0] === 'string'
+      ? { type: 'text', values: documents as string[] }
+      : { type: 'object', values: documents as JSONObject[] };
+
+  const baseTelemetryAttributes = getBaseTelemetryAttributes({
+    model,
+    telemetry,
+    headers,
+    settings: { maxRetries },
+  });
+
+  const tracer = getTracer(telemetry);
+
+  return recordSpan({
+    name: 'ai.rerank',
+    attributes: selectTelemetryAttributes({
+      telemetry,
+      attributes: {
+        ...assembleOperationName({ operationId: 'ai.rerank', telemetry }),
+        ...baseTelemetryAttributes,
+        'ai.documents': {
+          input: () => documents.map(document => JSON.stringify(document)),
+        },
+      },
+    }),
+    tracer,
+    fn: async () => {
+      const { ranking, response, providerMetadata, warnings } = await retry(
+        () =>
+          recordSpan({
+            name: 'ai.rerank.doRerank',
+            attributes: selectTelemetryAttributes({
+              telemetry,
+              attributes: {
+                ...assembleOperationName({
+                  operationId: 'ai.rerank.doRerank',
+                  telemetry,
+                }),
+                ...baseTelemetryAttributes,
+                // specific settings that only make sense on the outer level:
+                'ai.documents': {
+                  input: () =>
+                    documents.map(document => JSON.stringify(document)),
+                },
+              },
+            }),
+            tracer,
+            fn: async doRerankSpan => {
+              const modelResponse = await model.doRerank({
+                documents: documentsToSend,
+                query,
+                topN,
+                providerOptions,
+                abortSignal,
+                headers,
+              });
+
+              const ranking = modelResponse.ranking;
+
+              doRerankSpan.setAttributes(
+                await selectTelemetryAttributes({
+                  telemetry,
+                  attributes: {
+                    'ai.ranking.type': documentsToSend.type,
+                    'ai.ranking': {
+                      output: () =>
+                        ranking.map(ranking => JSON.stringify(ranking)),
+                    },
+                  },
+                }),
+              );
+
+              return {
+                ranking,
+                providerMetadata: modelResponse.providerMetadata,
+                response: modelResponse.response,
+                warnings: modelResponse.warnings,
+              };
+            },
+          }),
+      );
+
+      logWarnings({
+        warnings: warnings ?? [],
+        provider: model.provider,
+        model: model.modelId,
+      });
+
+      return new DefaultRerankResult({
+        originalDocuments: documents,
+        ranking: ranking.map(ranking => ({
+          originalIndex: ranking.index,
+          score: ranking.relevanceScore,
+          document: documents[ranking.index],
+        })),
+        providerMetadata,
+        response: {
+          id: response?.id,
+          timestamp: response?.timestamp ?? new Date(),
+          modelId: response?.modelId ?? model.modelId,
+          headers: response?.headers,
+          body: response?.body,
+        },
+      });
+    },
+  });
+}
+
+class DefaultRerankResult<VALUE> implements RerankResult<VALUE> {
+  readonly originalDocuments: RerankResult<VALUE>['originalDocuments'];
+  readonly ranking: RerankResult<VALUE>['ranking'];
+  readonly response: RerankResult<VALUE>['response'];
+  readonly providerMetadata: RerankResult<VALUE>['providerMetadata'];
+
+  constructor(options: {
+    originalDocuments: RerankResult<VALUE>['originalDocuments'];
+    ranking: RerankResult<VALUE>['ranking'];
+    providerMetadata?: RerankResult<VALUE>['providerMetadata'];
+    response: RerankResult<VALUE>['response'];
+  }) {
+    this.originalDocuments = options.originalDocuments;
+    this.ranking = options.ranking;
+    this.response = options.response;
+    this.providerMetadata = options.providerMetadata;
+  }
+
+  get rerankedDocuments(): RerankResult<VALUE>['rerankedDocuments'] {
+    return this.ranking.map(ranking => ranking.document);
+  }
+}

package/src/telemetry/assemble-operation-name.ts

@@ -0,0 +1,21 @@
+import { TelemetrySettings } from './telemetry-settings';
+
+export function assembleOperationName({
+  operationId,
+  telemetry,
+}: {
+  operationId: string;
+  telemetry?: TelemetrySettings;
+}) {
+  return {
+    // standardized operation and resource name:
+    'operation.name': `${operationId}${
+      telemetry?.functionId != null ? ` ${telemetry.functionId}` : ''
+    }`,
+    'resource.name': telemetry?.functionId,
+
+    // detailed, AI SDK specific data:
+    'ai.operationId': operationId,
+    'ai.telemetry.functionId': telemetry?.functionId,
+  };
+}

package/src/telemetry/get-base-telemetry-attributes.ts

@@ -0,0 +1,53 @@
+import { Attributes, AttributeValue } from '@opentelemetry/api';
+import { CallSettings, getTotalTimeoutMs } from '../prompt/call-settings';
+import { TelemetrySettings } from './telemetry-settings';
+
+export function getBaseTelemetryAttributes({
+  model,
+  settings,
+  telemetry,
+  headers,
+}: {
+  model: { modelId: string; provider: string };
+  settings: Omit<CallSettings, 'abortSignal' | 'headers' | 'temperature'>;
+  telemetry: TelemetrySettings | undefined;
+  headers: Record<string, string | undefined> | undefined;
+}): Attributes {
+  return {
+    'ai.model.provider': model.provider,
+    'ai.model.id': model.modelId,
+
+    // settings:
+    ...Object.entries(settings).reduce((attributes, [key, value]) => {
+      // Handle timeout specially since it can be a number or object
+      if (key === 'timeout') {
+        const totalTimeoutMs = getTotalTimeoutMs(
+          value as Parameters<typeof getTotalTimeoutMs>[0],
+        );
+        if (totalTimeoutMs != null) {
+          attributes[`ai.settings.${key}`] = totalTimeoutMs;
+        }
+      } else {
+        attributes[`ai.settings.${key}`] = value as AttributeValue;
+      }
+      return attributes;
+    }, {} as Attributes),
+
+    // add metadata as attributes:
+    ...Object.entries(telemetry?.metadata ?? {}).reduce(
+      (attributes, [key, value]) => {
+        attributes[`ai.telemetry.metadata.${key}`] = value;
+        return attributes;
+      },
+      {} as Attributes,
+    ),
+
+    // request headers
+    ...Object.entries(headers ?? {}).reduce((attributes, [key, value]) => {
+      if (value !== undefined) {
+        attributes[`ai.request.headers.${key}`] = value;
+      }
+      return attributes;
+    }, {} as Attributes),
+  };
+}

package/src/telemetry/get-tracer.ts

@@ -0,0 +1,20 @@
+import { Tracer, trace } from '@opentelemetry/api';
+import { noopTracer } from './noop-tracer';
+
+export function getTracer({
+  isEnabled = false,
+  tracer,
+}: {
+  isEnabled?: boolean;
+  tracer?: Tracer;
+} = {}): Tracer {
+  if (!isEnabled) {
+    return noopTracer;
+  }
+
+  if (tracer) {
+    return tracer;
+  }
+
+  return trace.getTracer('ai');
+}

package/src/telemetry/noop-tracer.ts

@@ -0,0 +1,69 @@
+import { Span, SpanContext, Tracer } from '@opentelemetry/api';
+
+/**
+ * Tracer implementation that does nothing (null object).
+ */
+export const noopTracer: Tracer = {
+  startSpan(): Span {
+    return noopSpan;
+  },
+
+  startActiveSpan<F extends (span: Span) => unknown>(
+    name: unknown,
+    arg1: unknown,
+    arg2?: unknown,
+    arg3?: F,
+  ): ReturnType<any> {
+    if (typeof arg1 === 'function') {
+      return arg1(noopSpan);
+    }
+    if (typeof arg2 === 'function') {
+      return arg2(noopSpan);
+    }
+    if (typeof arg3 === 'function') {
+      return arg3(noopSpan);
+    }
+  },
+};
+
+const noopSpan: Span = {
+  spanContext() {
+    return noopSpanContext;
+  },
+  setAttribute() {
+    return this;
+  },
+  setAttributes() {
+    return this;
+  },
+  addEvent() {
+    return this;
+  },
+  addLink() {
+    return this;
+  },
+  addLinks() {
+    return this;
+  },
+  setStatus() {
+    return this;
+  },
+  updateName() {
+    return this;
+  },
+  end() {
+    return this;
+  },
+  isRecording() {
+    return false;
+  },
+  recordException() {
+    return this;
+  },
+};
+
+const noopSpanContext: SpanContext = {
+  traceId: '',
+  spanId: '',
+  traceFlags: 0,
+};

package/src/telemetry/record-span.ts

@@ -0,0 +1,63 @@
+import { Attributes, Span, Tracer, SpanStatusCode } from '@opentelemetry/api';
+
+export async function recordSpan<T>({
+  name,
+  tracer,
+  attributes,
+  fn,
+  endWhenDone = true,
+}: {
+  name: string;
+  tracer: Tracer;
+  attributes: Attributes | PromiseLike<Attributes>;
+  fn: (span: Span) => Promise<T>;
+  endWhenDone?: boolean;
+}) {
+  return tracer.startActiveSpan(
+    name,
+    { attributes: await attributes },
+    async span => {
+      try {
+        const result = await fn(span);
+
+        if (endWhenDone) {
+          span.end();
+        }
+
+        return result;
+      } catch (error) {
+        try {
+          recordErrorOnSpan(span, error);
+        } finally {
+          // always stop the span when there is an error:
+          span.end();
+        }
+
+        throw error;
+      }
+    },
+  );
+}
+
+/**
+ * Record an error on a span. If the error is an instance of Error, an exception event will be recorded on the span, otherwise
+ * the span will be set to an error status.
+ *
+ * @param span - The span to record the error on.
+ * @param error - The error to record on the span.
+ */
+export function recordErrorOnSpan(span: Span, error: unknown) {
+  if (error instanceof Error) {
+    span.recordException({
+      name: error.name,
+      message: error.message,
+      stack: error.stack,
+    });
+    span.setStatus({
+      code: SpanStatusCode.ERROR,
+      message: error.message,
+    });
+  } else {
+    span.setStatus({ code: SpanStatusCode.ERROR });
+  }
+}

package/src/telemetry/select-telemetry-attributes.ts

@@ -0,0 +1,78 @@
+import type { Attributes, AttributeValue } from '@opentelemetry/api';
+import type { TelemetrySettings } from './telemetry-settings';
+
+type ResolvableAttributeValue = () =>
+  | AttributeValue
+  | PromiseLike<AttributeValue>
+  | undefined;
+
+export async function selectTelemetryAttributes({
+  telemetry,
+  attributes,
+}: {
+  telemetry?: TelemetrySettings;
+  attributes: {
+    [attributeKey: string]:
+      | AttributeValue
+      | { input: ResolvableAttributeValue }
+      | { output: ResolvableAttributeValue }
+      | undefined;
+  };
+}): Promise<Attributes> {
+  // when telemetry is disabled, return an empty object to avoid serialization overhead:
+  if (telemetry?.isEnabled !== true) {
+    return {};
+  }
+
+  const resultAttributes: Attributes = {};
+
+  for (const [key, value] of Object.entries(attributes)) {
+    if (value == null) {
+      continue;
+    }
+
+    // input value, check if it should be recorded:
+    if (
+      typeof value === 'object' &&
+      'input' in value &&
+      typeof value.input === 'function'
+    ) {
+      // default to true:
+      if (telemetry?.recordInputs === false) {
+        continue;
+      }
+
+      const result = await value.input();
+
+      if (result != null) {
+        resultAttributes[key] = result;
+      }
+
+      continue;
+    }
+
+    // output value, check if it should be recorded:
+    if (
+      typeof value === 'object' &&
+      'output' in value &&
+      typeof value.output === 'function'
+    ) {
+      // default to true:
+      if (telemetry?.recordOutputs === false) {
+        continue;
+      }
+
+      const result = await value.output();
+
+      if (result != null) {
+        resultAttributes[key] = result;
+      }
+      continue;
+    }
+
+    // value is an attribute value already:
+    resultAttributes[key] = value as AttributeValue;
+  }
+
+  return resultAttributes;
+}

package/src/telemetry/select-temetry-attributes.test.ts

@@ -0,0 +1,114 @@
+import { selectTelemetryAttributes } from './select-telemetry-attributes';
+import { it, expect } from 'vitest';
+
+it('should return an empty object when telemetry is disabled', async () => {
+  const result = await selectTelemetryAttributes({
+    telemetry: { isEnabled: false },
+    attributes: { key: 'value' },
+  });
+  expect(result).toEqual({});
+});
+
+it('should return an empty object when telemetry enablement is undefined', async () => {
+  const result = await selectTelemetryAttributes({
+    telemetry: { isEnabled: undefined },
+    attributes: { key: 'value' },
+  });
+  expect(result).toEqual({});
+});
+
+it('should return attributes with simple values', async () => {
+  const result = await selectTelemetryAttributes({
+    telemetry: { isEnabled: true },
+    attributes: { string: 'value', number: 42, boolean: true },
+  });
+  expect(result).toEqual({ string: 'value', number: 42, boolean: true });
+});
+
+it('should handle input functions when recordInputs is true', async () => {
+  const result = await selectTelemetryAttributes({
+    telemetry: { isEnabled: true, recordInputs: true },
+    attributes: {
+      input: { input: () => 'input value' },
+      other: 'other value',
+    },
+  });
+  expect(result).toEqual({ input: 'input value', other: 'other value' });
+});
+
+it('should not include input functions when recordInputs is false', async () => {
+  const result = await selectTelemetryAttributes({
+    telemetry: { isEnabled: true, recordInputs: false },
+    attributes: {
+      input: { input: () => 'input value' },
+      other: 'other value',
+    },
+  });
+  expect(result).toEqual({ other: 'other value' });
+});
+
+it('should handle output functions when recordOutputs is true', async () => {
+  const result = await selectTelemetryAttributes({
+    telemetry: { isEnabled: true, recordOutputs: true },
+    attributes: {
+      output: { output: () => 'output value' },
+      other: 'other value',
+    },
+  });
+  expect(result).toEqual({ output: 'output value', other: 'other value' });
+});
+
+it('should not include output functions when recordOutputs is false', async () => {
+  const result = await selectTelemetryAttributes({
+    telemetry: { isEnabled: true, recordOutputs: false },
+    attributes: {
+      output: { output: () => 'output value' },
+      other: 'other value',
+    },
+  });
+  expect(result).toEqual({ other: 'other value' });
+});
+
+it('should ignore undefined values', async () => {
+  const result = await selectTelemetryAttributes({
+    telemetry: { isEnabled: true },
+    attributes: {
+      defined: 'value',
+      undefined: undefined,
+    },
+  });
+  expect(result).toEqual({ defined: 'value' });
+});
+
+it('should ignore input and output functions that return undefined', async () => {
+  const result = await selectTelemetryAttributes({
+    telemetry: { isEnabled: true },
+    attributes: {
+      input: { input: () => undefined },
+      output: { output: () => undefined },
+      other: 'value',
+    },
+  });
+  expect(result).toEqual({ other: 'value' });
+});
+
+it('should handle mixed attribute types correctly', async () => {
+  const result = await selectTelemetryAttributes({
+    telemetry: { isEnabled: true },
+    attributes: {
+      simple: 'value',
+      input: { input: () => 'input value' },
+      output: { output: () => 'output value' },
+      undefined: undefined,
+      // Invalid null input
+      null: null as any,
+      input_null: { input: () => null as any },
+    },
+  });
+
+  expect(result).toEqual({
+    simple: 'value',
+    input: 'input value',
+    output: 'output value',
+  });
+});