@tanstack/ai 0.43.0 → 0.44.0

package/src/activities/rerank/index.ts

@@ -0,0 +1,302 @@
+/**
+ * Rerank Activity
+ *
+ * Reorders a set of documents by semantic relevance to a query.
+ * This is a self-contained module with implementation, types, and JSDoc.
+ */
+
+import { aiEventClient } from '@tanstack/ai-event-client'
+import { resolveDebugOption } from '../../logger/resolve'
+import { isAbortShapedError } from '../error-payload'
+import {
+  createGenerationContext,
+  runGenerationAbort,
+  runGenerationError,
+  runGenerationFinish,
+  runGenerationStart,
+  runGenerationUsage,
+} from '../middleware/run'
+import type { InternalLogger } from '../../logger/internal-logger'
+import type { DebugOption } from '../../logger/types'
+import type { GenerationMiddleware } from '../middleware/types'
+import type { RerankAdapter } from './adapter'
+import type { RerankResult } from '../../types'
+
+// ===========================
+// Activity Kind
+// ===========================
+
+/** The adapter kind this activity handles */
+export const kind = 'rerank' as const
+
+// ===========================
+// Type Extraction Helpers
+// ===========================
+
+/** Extract provider options from a RerankAdapter via ~types */
+export type RerankProviderOptions<TAdapter> = TAdapter extends {
+  '~types': { providerOptions: infer P extends object }
+}
+  ? P
+  : object
+
+// ===========================
+// Activity Options Type
+// ===========================
+
+/**
+ * Options for the rerank activity. The model is extracted from the adapter's
+ * model property.
+ *
+ * @template TAdapter - The rerank adapter type
+ * @template TDocument - The document element type (string or object)
+ */
+export interface RerankActivityOptions<
+  TAdapter extends RerankAdapter<string, RerankProviderOptions<TAdapter>>,
+  TDocument extends string | object = string,
+> {
+  /** The rerank adapter to use (must be created with a model) */
+  adapter: TAdapter & { kind: typeof kind }
+  /** The query documents are scored against. */
+  query: string
+  /**
+   * Documents to rerank. Either strings or JSON-serializable objects — object
+   * documents are serialized with `JSON.stringify` before being sent to the
+   * provider, and the original element (string or object) is returned in the
+   * result, preserving its type.
+   */
+  documents: Array<TDocument>
+  /** Return only the top N results. */
+  topN?: number
+  /** Provider-specific options */
+  modelOptions?: RerankProviderOptions<TAdapter>
+  /** Forwarded to the provider request for cancellation. */
+  abortSignal?: AbortSignal
+  /**
+   * Observe-only middleware notified on start, usage, success, abort, and
+   * error. Pass `otelMiddleware()` to emit OpenTelemetry spans, or implement
+   * the `GenerationMiddleware` contract for a custom backend.
+   */
+  middleware?: Array<GenerationMiddleware>
+  /**
+   * Enable debug logging. Pass `true` to enable all categories, `false` to
+   * silence everything including errors, or a `DebugConfig` object for granular
+   * control and/or a custom `Logger`.
+   */
+  debug?: DebugOption
+}
+
+// ===========================
+// Helper Functions
+// ===========================
+
+function createId(prefix: string): string {
+  return `${prefix}-${Date.now()}-${Math.random().toString(36).slice(2, 9)}`
+}
+
+/** Serialize a document for the provider. Strings pass through untouched. */
+function serializeDocument(document: string | object): string {
+  return typeof document === 'string' ? document : JSON.stringify(document)
+}
+
+function isAbortError(error: unknown, signal?: AbortSignal): boolean {
+  // Prefer the error's own identity over the signal state. A genuine
+  // cancellation throws an abort-shaped error (DOM `AbortError`, the OpenRouter
+  // SDK's `RequestAbortedError`, …). Classifying on `signal.aborted` alone would
+  // misroute a real failure — e.g. the out-of-range-index throw below — to the
+  // abort hook whenever a shared/long-lived signal happens to already be
+  // aborted, hiding it from `onError` observers.
+  if (isAbortShapedError(error)) return true
+  // Fall back to signal state only for non-Error throws we can't otherwise
+  // identify; a real Error with a non-abort name is never an abort.
+  return error instanceof Error ? false : signal?.aborted === true
+}
+
+// ===========================
+// Activity Implementation
+// ===========================
+
+/**
+ * Rerank activity - reorders documents by relevance to a query.
+ *
+ * @example Basic reranking
+ * ```ts
+ * import { rerank } from '@tanstack/ai'
+ * import { cohereRerank } from '@tanstack/ai-cohere'
+ *
+ * const { ranking, rerankedDocuments } = await rerank({
+ *   adapter: cohereRerank('rerank-v3.5'),
+ *   query: 'talk about rain',
+ *   documents: ['sunny day at the beach', 'rainy afternoon in the city'],
+ *   topN: 2,
+ * })
+ *
+ * console.log(rerankedDocuments[0]) // 'rainy afternoon in the city'
+ * ```
+ *
+ * @example Reranking object documents
+ * ```ts
+ * const { ranking } = await rerank({
+ *   adapter: cohereRerank('rerank-v3.5'),
+ *   query: 'best laptop for travel',
+ *   documents: [
+ *     { id: 1, text: 'A heavy gaming desktop' },
+ *     { id: 2, text: 'A lightweight ultrabook with all-day battery' },
+ *   ],
+ * })
+ *
+ * // ranking[0].document is the original object, fully typed.
+ * console.log(ranking[0].document.id)
+ * ```
+ */
+export async function rerank<
+  TAdapter extends RerankAdapter<string, RerankProviderOptions<TAdapter>>,
+  TDocument extends string | object = string,
+>(
+  options: RerankActivityOptions<TAdapter, TDocument>,
+): Promise<RerankResult<TDocument>> {
+  const {
+    adapter,
+    query,
+    documents,
+    topN,
+    modelOptions,
+    abortSignal,
+    middleware,
+  } = options
+  const model = adapter.model
+  const requestId = createId('rerank')
+  const startTime = Date.now()
+  const logger: InternalLogger = resolveDebugOption(options.debug)
+
+  if (documents.length === 0) {
+    throw new Error('rerank() requires at least one document')
+  }
+
+  const mwCtx = createGenerationContext({
+    requestId,
+    // `rerank` joins the GenerationActivity union; otel maps it to its own
+    // gen_ai.operation.name.
+    activity: 'rerank',
+    provider: adapter.name,
+    model,
+    modelOptions,
+    createId,
+  })
+
+  await runGenerationStart(middleware, mwCtx)
+
+  aiEventClient.emit('rerank:request:started', {
+    requestId,
+    provider: adapter.name,
+    model,
+    documentCount: documents.length,
+    timestamp: startTime,
+  })
+
+  logger.request(`activity=rerank provider=${adapter.name}`, {
+    provider: adapter.name,
+    model,
+    documentCount: documents.length,
+  })
+
+  // Serialize once; reuse for the request only. Original documents are mapped
+  // back by index below so the caller's element type is preserved.
+  const serialized = documents.map(serializeDocument)
+
+  try {
+    const result = await adapter.rerank({
+      model,
+      query,
+      documents: serialized,
+      topN,
+      modelOptions,
+      abortSignal,
+      logger,
+    })
+
+    const ranking = result.ranking.map((r) => {
+      const document = documents[r.index]
+      if (document === undefined) {
+        throw new Error(
+          `rerank(): provider ${adapter.name} returned out-of-range index ${r.index}`,
+        )
+      }
+      return { index: r.index, score: r.score, document }
+    })
+    const rerankedDocuments = ranking.map((r) => r.document)
+
+    const duration = Date.now() - startTime
+
+    aiEventClient.emit('rerank:request:completed', {
+      requestId,
+      provider: adapter.name,
+      model,
+      documentCount: documents.length,
+      resultCount: ranking.length,
+      duration,
+      timestamp: Date.now(),
+    })
+
+    aiEventClient.emit('rerank:usage', {
+      requestId,
+      model,
+      usage: result.usage,
+      timestamp: Date.now(),
+    })
+
+    logger.output(`activity=rerank results=${ranking.length}`, {
+      resultCount: ranking.length,
+    })
+
+    await runGenerationUsage(middleware, mwCtx, result.usage)
+    await runGenerationFinish(middleware, mwCtx, {
+      duration,
+      usage: result.usage,
+    })
+
+    return {
+      id: result.id,
+      model,
+      ranking,
+      rerankedDocuments,
+      usage: result.usage,
+    }
+  } catch (error) {
+    const duration = Date.now() - startTime
+    if (isAbortError(error, abortSignal)) {
+      await runGenerationAbort(middleware, mwCtx, {
+        reason: error instanceof Error ? error.message : undefined,
+        duration,
+      })
+    } else {
+      await runGenerationError(middleware, mwCtx, { error, duration })
+    }
+    logger.errors('rerank activity failed', { error, source: 'rerank' })
+    throw error
+  }
+}
+
+// ===========================
+// Options Factory
+// ===========================
+
+/**
+ * Create typed options for the rerank() function without executing.
+ */
+export function createRerankOptions<
+  TAdapter extends RerankAdapter<string, RerankProviderOptions<TAdapter>>,
+  TDocument extends string | object = string,
+>(
+  options: RerankActivityOptions<TAdapter, TDocument>,
+): RerankActivityOptions<TAdapter, TDocument> {
+  return options
+}
+
+// Re-export adapter types
+export type {
+  RerankAdapter,
+  RerankAdapterConfig,
+  AnyRerankAdapter,
+} from './adapter'
+export { BaseRerankAdapter } from './adapter'

package/src/activities/summarize/index.ts

@@ -17,6 +17,12 @@ import {
   runGenerationStart,
   runGenerationUsage,
 } from '../middleware/run'
+import {
+  abortReasonMessage,
+  createActivityAbortControls,
+  isActivityAbortError,
+  raceWithAbort,
+} from '../../utilities/activity-abort'
 import type { InternalLogger } from '../../logger/internal-logger'
 import type { DebugOption } from '../../logger/types'
 import type { GenerationMiddleware } from '../middleware/types'
@@ -35,10 +41,11 @@ export const kind = 'summarize' as const
 // ===========================
 
 /** Extract provider options from a SummarizeAdapter via ~types */
-export type SummarizeProviderOptions<TAdapter> =
-  TAdapter extends SummarizeAdapter<any, any>
-    ? TAdapter['~types']['providerOptions']
-    : object
+export type SummarizeProviderOptions<TAdapter> = TAdapter extends {
+  '~types': { providerOptions: infer P extends object }
+}
+  ? P
+  : object
 
 // ===========================
 // Activity Options Type
@@ -93,6 +100,18 @@ export interface SummarizeActivityOptions<
    * mid-summary fires `onAbort`.
    */
   middleware?: Array<GenerationMiddleware>
+  /**
+   * Maximum duration of this activity invocation in milliseconds.
+   * No SDK-wide default — choose a value suitable for the provider and job.
+   * Composed with {@link abortSignal}; the first abort wins.
+   */
+  timeout?: number
+  /**
+   * Caller cancellation signal (request disconnects, job/runtime cancellation).
+   * Composed with {@link timeout} into an effective signal forwarded to the
+   * adapter. Request-specific — not stored on global provider client config.
+   */
+  abortSignal?: AbortSignal
   /**
    * Whether to stream the summarization result.
    * When true, returns an AsyncIterable<StreamChunk> for streaming output.
@@ -195,18 +214,12 @@ export function summarize<
 
   if (stream) {
     return runStreamingSummarize(
-      options as SummarizeActivityOptions<
-        SummarizeAdapter<string, object>,
-        true
-      >,
+      options as SummarizeActivityOptions<TAdapter, true>,
     ) as SummarizeActivityResult<TStream>
   }
 
   return runSummarize(
-    options as SummarizeActivityOptions<
-      SummarizeAdapter<string, object>,
-      false
-    >,
+    options as SummarizeActivityOptions<TAdapter, false>,
   ) as SummarizeActivityResult<TStream>
 }
 
@@ -216,13 +229,26 @@ export function summarize<
 async function runSummarize(
   options: SummarizeActivityOptions<SummarizeAdapter<string, object>, false>,
 ): Promise<SummarizationResult> {
-  const { adapter, text, maxLength, style, focus, modelOptions, middleware } =
-    options
+  const {
+    adapter,
+    text,
+    maxLength,
+    style,
+    focus,
+    modelOptions,
+    middleware,
+    timeout,
+    abortSignal: callerAbortSignal,
+  } = options
   const model = adapter.model
   const requestId = createId('summarize')
   const inputLength = text.length
   const startTime = Date.now()
   const logger: InternalLogger = resolveDebugOption(options.debug)
+  const abortControls = createActivityAbortControls({
+    timeout,
+    abortSignal: callerAbortSignal,
+  })
 
   const mwCtx = createGenerationContext({
     requestId,
@@ -259,10 +285,15 @@ async function runSummarize(
     focus,
     modelOptions,
     logger,
+    ...(abortControls.signal ? { abortSignal: abortControls.signal } : {}),
   }
 
   try {
-    const rawResult = await adapter.summarize(summarizeOptions)
+    const rawResult = await raceWithAbort(
+      adapter.summarize(summarizeOptions),
+      abortControls.signal,
+    )
+    abortControls.clear()
     // Transforms run before anything observes the result — the same order every
     // media activity uses — so the run record and the returned value are the
     // same object.
@@ -294,10 +325,19 @@ async function runSummarize(
 
     return result
   } catch (error) {
-    await runGenerationError(middleware, mwCtx, {
-      error,
-      duration: Date.now() - startTime,
-    })
+    abortControls.clear()
+    const duration = Date.now() - startTime
+    if (isActivityAbortError(error, abortControls.signal)) {
+      await runGenerationAbort(middleware, mwCtx, {
+        reason: abortReasonMessage(error, abortControls.signal),
+        duration,
+      })
+    } else {
+      await runGenerationError(middleware, mwCtx, {
+        error,
+        duration,
+      })
+    }
     logger.errors('summarize activity failed', {
       error,
       source: 'summarize',

package/src/index.ts

@@ -2,22 +2,26 @@
 export {
   chat,
   summarize,
+  rerank,
   generateImage,
   generateAudio,
   generateVideo,
   getVideoJobStatus,
   generateSpeech,
   generateTranscription,
+  embed,
 } from './activities/index'
 
 // Create options functions - for pre-defining typed configurations
 export { createChatOptions } from './activities/chat/index'
 export { createSummarizeOptions } from './activities/summarize/index'
+export { createRerankOptions } from './activities/rerank/index'
 export { createImageOptions } from './activities/generateImage/index'
 export { createAudioOptions } from './activities/generateAudio/index'
 export { createVideoOptions } from './activities/generateVideo/index'
 export { createSpeechOptions } from './activities/generateSpeech/index'
 export { createTranscriptionOptions } from './activities/generateTranscription/index'
+export { createEmbedOptions } from './activities/embed/index'
 
 // Re-export types
 export type {
@@ -36,8 +40,15 @@ export type {
   TranscriptionAdapter,
   AnyVideoAdapter,
   VideoAdapter,
+  AnyEmbeddingAdapter,
+  EmbeddingAdapter,
+  AnyRerankAdapter,
+  RerankAdapter,
 } from './activities/index'
 
+// Rerank adapter base + types
+export { BaseRerankAdapter } from './activities/rerank/adapter'
+
 // Tool definition
 export {
   toolDefinition,
@@ -303,6 +314,14 @@ export { buildBaseUsage, type BaseUsageInput } from './utilities/usage'
 export { resolveMediaPrompt } from './utilities/media-prompt'
 export type { ResolvedMediaPrompt } from './utilities/media-prompt'
 
+// Embedding input resolution (used by embedding adapters)
+export {
+  resolveEmbeddingInput,
+  requireTextOnlyEmbeddingInput,
+  countEmbeddingInputModalities,
+} from './utilities/embedding-input'
+export type { ResolvedEmbeddingItem } from './utilities/embedding-input'
+
 // System prompts (type + normaliser used by adapters)
 export type { SystemPrompt, NormalizedSystemPrompt } from './system-prompts'
 export { normalizeSystemPrompts } from './system-prompts'

package/src/middlewares/otel.ts

@@ -28,12 +28,15 @@ import type {
   GenerationMiddleware,
   GenerationMiddlewareContext,
 } from '../activities/middleware/types'
+import type { TokenUsage } from '../types'
 
 /**
  * Scope (role) of an OTel span emitted by this middleware.
  *
  * - `chat` — the root span for a single `chat()` call
- * - `iteration` — one per agent-loop iteration (one model call)
+ * - `iteration` — one per provider model call (agent-loop `beforeModel`
+ *   turn, or the separate `structuredOutput` finalization when
+ *   `outputSchema` skips the agent loop — see #1054)
  * - `tool` — one per tool execution inside an iteration
  * - `generation` — the single span for a media activity call
  *   (`generateImage`, `generateVideo`, `generateSpeech`, …)
@@ -83,6 +86,8 @@ const OPERATION_NAME: Record<GenerationActivity, string> = {
   audio: 'audio_generation',
   tts: 'text_to_speech',
   transcription: 'transcription',
+  embedding: 'embeddings',
+  rerank: 'rerank',
   summarize: 'summarize',
 }
 
@@ -139,6 +144,8 @@ interface RequestState {
    * from the (base-shaped) finish info, which doesn't carry it.
    */
   lastFinishReason: string | null
+  rootUsageAttributes: Record<string, number> | null
+  rootUsageApplied: boolean
 }
 
 const stateByCtx = new WeakMap<ChatMiddlewareContext, RequestState>()
@@ -146,6 +153,29 @@ const stateByCtx = new WeakMap<ChatMiddlewareContext, RequestState>()
 const DEFAULT_MAX_CONTENT_LENGTH = 100_000
 const REDACTION_FAILED_SENTINEL = '[redaction_failed]'
 
+function accumulateUsageAttributes(
+  current: Record<string, number> | null,
+  usage: TokenUsage,
+): Record<string, number> {
+  const accumulated = current ?? {}
+  for (const [key, value] of Object.entries(usageAttributes(usage))) {
+    if (typeof value === 'number') {
+      accumulated[key] = (accumulated[key] ?? 0) + value
+    }
+  }
+  return accumulated
+}
+
+function applyRootUsage(state: RequestState, fallbackUsage?: TokenUsage): void {
+  if (state.rootUsageApplied) return
+
+  const attributes =
+    state.rootUsageAttributes ??
+    (fallbackUsage ? usageAttributes(fallbackUsage) : null)
+  if (attributes) state.rootSpan.setAttributes(attributes)
+  state.rootUsageApplied = true
+}
+
 function serializeContent(content: unknown): string {
   if (typeof content === 'string') return content
   if (!Array.isArray(content)) return ''
@@ -396,12 +426,23 @@ export function otelMiddleware(
           assistantTextBufferTruncated: false,
           startTime: Date.now(),
           lastFinishReason: null,
+          rootUsageAttributes: null,
+          rootUsageApplied: false,
         })
       })
     },
 
     onConfig(ctx, config) {
-      if (ctx.phase !== 'beforeModel') return
+      // Open an iteration span for every provider model call:
+      // - `beforeModel`: agent-loop chatStream turns
+      // - `structuredOutput`: separate structured-output finalization
+      //   (no-tools + outputSchema skips the agent loop, so without this
+      //   phase there is no generation span and captureContent is a silent
+      //   no-op — see #1054).
+      // Native-combined mode never fires `structuredOutput`, so a run that
+      // already opened spans via `beforeModel` is not double-counted.
+      if (ctx.phase !== 'beforeModel' && ctx.phase !== 'structuredOutput')
+        return
       safeCall('otel.onConfig', () => {
         const state = stateByCtx.get(ctx)
         if (!state) return
@@ -411,20 +452,26 @@ export function otelMiddleware(
         // on it. Close it here, just before opening the next iteration.
         closeIterationSpan(state, ctx)
 
+        // Number spans by the order of model calls this middleware has seen,
+        // not by `ctx.iteration`. After an agent-loop turn, structured-output
+        // finalization reuses the engine's last iteration index; using our
+        // own counter keeps finalization as a distinct leaf (#N+1).
+        const iteration = state.iterationCount
+
         const info: OtelSpanInfo<'iteration'> = {
           kind: 'iteration',
           ctx,
-          iteration: ctx.iteration,
+          iteration,
         }
         const name =
           safeCall('otel.spanNameFormatter', () => spanNameFormatter?.(info)) ??
-          `chat ${ctx.model} #${ctx.iteration}`
+          `chat ${ctx.model} #${iteration}`
 
         const baseAttrs: Record<string, AttributeValue> = {
           'gen_ai.system': ctx.provider,
           'gen_ai.operation.name': 'chat',
           'gen_ai.request.model': ctx.model,
-          'tanstack.ai.iteration': ctx.iteration,
+          'tanstack.ai.iteration': iteration,
         }
         // Sampling options now live in provider-native `modelOptions`, and
         // providers spell them differently (e.g. `max_output_tokens`,
@@ -651,6 +698,11 @@ export function otelMiddleware(
         const state = stateByCtx.get(chatCtx)
         if (!state) return
 
+        state.rootUsageAttributes = accumulateUsageAttributes(
+          state.rootUsageAttributes,
+          usage,
+        )
+
         // Always record the token histogram — metrics don't depend on having
         // an iteration span, and skipping here would drop metric data if an
         // adapter emits `onUsage` outside the iteration window.
@@ -890,6 +942,7 @@ export function otelMiddleware(
           })
         }
 
+        applyRootUsage(state)
         safeCall('otel.onSpanEnd', () =>
           onSpanEnd?.({ kind: 'chat', ctx: chatCtx }, state.rootSpan),
         )
@@ -971,6 +1024,7 @@ export function otelMiddleware(
           })
         }
 
+        applyRootUsage(state)
         safeCall('otel.onSpanEnd', () =>
           onSpanEnd?.({ kind: 'chat', ctx: chatCtx }, state.rootSpan),
         )
@@ -1028,9 +1082,6 @@ export function otelMiddleware(
           })
         }
 
-        if (info.usage) {
-          state.rootSpan.setAttributes(usageAttributes(info.usage))
-        }
         if (state.lastFinishReason) {
           state.rootSpan.setAttribute('gen_ai.response.finish_reasons', [
             state.lastFinishReason,
@@ -1041,6 +1092,7 @@ export function otelMiddleware(
           state.iterationCount,
         )
 
+        applyRootUsage(state, info.usage)
         safeCall('otel.onSpanEnd', () =>
           onSpanEnd?.({ kind: 'chat', ctx: chatCtx }, state.rootSpan),
         )