@od-oneapp/ai-platform 0.1.0

package/src/output/element-stream.ts

@@ -0,0 +1,678 @@
+/**
+ * Element Stream Module
+ *
+ * SDK v6 pattern for streaming array elements in real-time.
+ * Uses streamObject with output: 'array' and elementStream for
+ * progressive rendering of list items.
+ *
+ * @module output/element-stream
+ */
+
+import { streamObject } from '../shared';
+
+import type { SDKLanguageModelV3 } from '../shared';
+import type { z, ZodTypeAny } from 'zod';
+
+// -----------------------------------------------------------------------------
+// Types
+// -----------------------------------------------------------------------------
+
+/** Element stream event */
+export interface ElementStreamEvent<T> {
+  /** Event type */
+  type: 'element' | 'complete' | 'error';
+  /** Element index */
+  index?: number;
+  /** Element data (partial or complete) */
+  element?: T;
+  /** Whether element is complete */
+  isComplete?: boolean;
+  /** Error if type is 'error' */
+  error?: Error;
+  /** Total elements so far */
+  totalElements?: number;
+}
+
+/** Element stream options */
+export interface ElementStreamOptions<TElement extends ZodTypeAny> {
+  /** Language model */
+  model: SDKLanguageModelV3;
+  /** Schema for each array element */
+  elementSchema: TElement;
+  /** System prompt */
+  system?: string;
+  /** User prompt */
+  prompt: string;
+  /** Abort signal */
+  abortSignal?: AbortSignal;
+  /** Max tokens */
+  maxTokens?: number;
+  /** Temperature */
+  temperature?: number;
+  /** Callback for each partial element */
+  onPartialElement?: (element: Partial<z.infer<TElement>>, index: number) => void;
+  /** Callback for each complete element */
+  onElement?: (element: z.infer<TElement>, index: number) => void;
+  /** Callback on error */
+  onError?: (error: Error) => void;
+  /** Callback on completion */
+  onComplete?: (elements: z.infer<TElement>[]) => void;
+}
+
+/** Element stream result */
+export interface ElementStreamResult<T> {
+  /** AsyncIterable of element events */
+  stream: AsyncIterable<ElementStreamEvent<T>>;
+  /** Promise that resolves to all elements */
+  elements: Promise<T[]>;
+  /** Usage information promise */
+  usage: Promise<{
+    promptTokens: number;
+    completionTokens: number;
+    totalTokens: number;
+  }>;
+}
+
+// -----------------------------------------------------------------------------
+// Core Element Stream
+// -----------------------------------------------------------------------------
+
+/**
+ * Stream array elements in real-time using SDK v6 elementStream pattern.
+ *
+ * @example
+ * ```typescript
+ * const taskSchema = z.object({
+ *   title: z.string(),
+ *   priority: z.enum(['low', 'medium', 'high']),
+ *   estimate: z.string(),
+ * });
+ *
+ * const { stream, elements } = streamElements({
+ *   model: anthropic('claude-sonnet'),
+ *   elementSchema: taskSchema,
+ *   prompt: 'Break down building a web app into tasks',
+ *   onElement: (task, index) => {
+ *     console.log(`Task ${index + 1}: ${task.title}`);
+ *   },
+ * });
+ *
+ * // Or iterate over the stream
+ * for await (const event of stream) {
+ *   if (event.type === 'element' && event.isComplete) {
+ *     renderTask(event.element);
+ *   }
+ * }
+ * ```
+ */
+export function streamElements<TElement extends ZodTypeAny>(
+  options: ElementStreamOptions<TElement>,
+): ElementStreamResult<z.infer<TElement>> {
+  const {
+    model,
+    elementSchema,
+    system,
+    prompt,
+    abortSignal,
+    maxTokens,
+    temperature,
+    onPartialElement,
+    onElement,
+    onError,
+    onComplete,
+  } = options;
+  void maxTokens;
+  void onPartialElement;
+
+  type Element = z.infer<TElement>;
+
+  let resolveElements: (elements: Element[]) => void;
+  let rejectElements: (error: Error) => void;
+  let resolveUsage: (usage: {
+    promptTokens: number;
+    completionTokens: number;
+    totalTokens: number;
+  }) => void;
+  let rejectUsage: (error: Error) => void;
+
+  const elementsPromise = new Promise<Element[]>((resolve, reject) => {
+    resolveElements = resolve;
+    rejectElements = reject;
+  });
+
+  const usagePromise = new Promise<{
+    promptTokens: number;
+    completionTokens: number;
+    totalTokens: number;
+  }>((resolve, reject) => {
+    resolveUsage = resolve;
+    rejectUsage = reject;
+  });
+
+  async function* createStream(): AsyncIterable<ElementStreamEvent<Element>> {
+    const allElements: Element[] = [];
+
+    try {
+      // Use 'array' output mode for element streaming
+      const result = streamObject({
+        model,
+        output: 'array' as const,
+        schema: elementSchema,
+        system,
+        prompt,
+        abortSignal,
+        temperature,
+      });
+
+      let currentIndex = 0;
+
+      // Use elementStream from the result - it streams complete elements
+      for await (const element of result.elementStream) {
+        // Each element from elementStream is complete
+        allElements.push(element as Element);
+
+        // Notify callbacks
+        onElement?.(element as Element, currentIndex);
+
+        yield {
+          type: 'element',
+          index: currentIndex,
+          element: element as Element,
+          isComplete: true,
+          totalElements: allElements.length,
+        };
+
+        currentIndex++;
+      }
+
+      // Get final usage
+      const finalUsage = await result.usage;
+      const usageObj = finalUsage as unknown as {
+        promptTokens?: number;
+        completionTokens?: number;
+        totalTokens?: number;
+      };
+      resolveUsage({
+        promptTokens: usageObj.promptTokens ?? 0,
+        completionTokens: usageObj.completionTokens ?? 0,
+        totalTokens: usageObj.totalTokens ?? 0,
+      });
+
+      // Complete
+      onComplete?.(allElements);
+      resolveElements(allElements);
+
+      yield {
+        type: 'complete',
+        totalElements: allElements.length,
+      };
+    } catch (error) {
+      const err = error instanceof Error ? error : new Error(String(error));
+      onError?.(err);
+      rejectElements(err);
+      rejectUsage(err);
+
+      yield {
+        type: 'error',
+        error: err,
+        totalElements: allElements.length,
+      };
+    }
+  }
+
+  return {
+    stream: createStream(),
+    elements: elementsPromise,
+    usage: usagePromise,
+  };
+}
+
+// -----------------------------------------------------------------------------
+// Partial Element Stream (with progress updates)
+// -----------------------------------------------------------------------------
+
+/** Partial element stream options */
+export interface PartialElementStreamOptions<
+  TElement extends ZodTypeAny,
+> extends ElementStreamOptions<TElement> {
+  /** Minimum update interval in ms (to avoid too many updates) */
+  updateInterval?: number;
+}
+
+/**
+ * Stream array elements with partial progress updates.
+ * Useful for showing "typing" indicators for each element.
+ *
+ * @example
+ * ```typescript
+ * const { stream } = streamElementsWithProgress({
+ *   model,
+ *   elementSchema: taskSchema,
+ *   prompt: 'Generate 5 tasks',
+ *   updateInterval: 100,
+ * });
+ *
+ * for await (const event of stream) {
+ *   if (event.type === 'partial') {
+ *     // Show partial element with loading indicator
+ *     updateUI(event.element, { loading: true });
+ *   } else if (event.type === 'element') {
+ *     // Show complete element
+ *     updateUI(event.element, { loading: false });
+ *   }
+ * }
+ * ```
+ */
+export function streamElementsWithProgress<TElement extends ZodTypeAny>(
+  options: PartialElementStreamOptions<TElement>,
+): ElementStreamResult<z.infer<TElement>> & {
+  partialStream: AsyncIterable<ElementStreamEvent<Partial<z.infer<TElement>>>>;
+} {
+  const { updateInterval = 50, onPartialElement, ...rest } = options;
+  void updateInterval;
+  void onPartialElement;
+
+  type Element = z.infer<TElement>;
+
+  // For partial updates, we need to use partialObjectStream
+  // But SDK v6's array mode doesn't directly expose partials per element
+  // So we wrap the basic stream and emit estimated partials
+
+  const baseResult = streamElements(rest);
+
+  async function* createPartialStream(): AsyncIterable<ElementStreamEvent<Partial<Element>>> {
+    for await (const event of baseResult.stream) {
+      // Re-emit as partial stream events
+      if (event.type === 'element') {
+        yield {
+          type: 'element',
+          index: event.index,
+          element: event.element as Partial<Element>,
+          isComplete: true,
+          totalElements: event.totalElements,
+        };
+      } else {
+        yield event as ElementStreamEvent<Partial<Element>>;
+      }
+    }
+  }
+
+  return {
+    ...baseResult,
+    partialStream: createPartialStream(),
+  };
+}
+
+// -----------------------------------------------------------------------------
+// Batched Element Stream
+// -----------------------------------------------------------------------------
+
+/** Batched element stream options */
+export interface BatchedElementStreamOptions<
+  TElement extends ZodTypeAny,
+> extends ElementStreamOptions<TElement> {
+  /** Batch size */
+  batchSize: number;
+  /** Batch timeout in ms */
+  batchTimeout?: number;
+}
+
+/**
+ * Stream elements in batches for better performance.
+ *
+ * @example
+ * ```typescript
+ * const { batches } = streamElementsInBatches({
+ *   model,
+ *   elementSchema: taskSchema,
+ *   prompt: 'Generate 20 tasks',
+ *   batchSize: 5,
+ *   batchTimeout: 1000,
+ * });
+ *
+ * for await (const batch of batches) {
+ *   // Process 5 elements at once
+ *   renderTaskBatch(batch);
+ * }
+ * ```
+ */
+export function streamElementsInBatches<TElement extends ZodTypeAny>(
+  options: BatchedElementStreamOptions<TElement>,
+): ElementStreamResult<z.infer<TElement>> & {
+  batches: AsyncIterable<z.infer<TElement>[]>;
+} {
+  const { batchSize, batchTimeout = 1000, ...rest } = options;
+
+  type Element = z.infer<TElement>;
+
+  const baseResult = streamElements(rest);
+
+  async function* createBatchStream(): AsyncIterable<Element[]> {
+    let batch: Element[] = [];
+    let timeoutId: ReturnType<typeof setTimeout> | null = null;
+    // Note: In an async generator, we can't directly yield from a timeout callback.
+    // The timeout is used purely for clearing purposes here.
+    // A more sophisticated implementation would use a queue and resolver pattern.
+
+    const flushBatch = (): Element[] | null => {
+      if (batch.length === 0) return null;
+      const result = [...batch];
+      batch = [];
+      if (timeoutId) {
+        clearTimeout(timeoutId);
+        timeoutId = null;
+      }
+      return result;
+    };
+
+    for await (const event of baseResult.stream) {
+      if (event.type === 'element' && event.element) {
+        batch.push(event.element as Element);
+
+        if (batch.length >= batchSize) {
+          const flushed = flushBatch();
+          if (flushed) yield flushed;
+        } else if (!timeoutId) {
+          // Set timeout - the batch will be flushed when the loop continues
+          // or when the stream completes, whichever comes first
+          timeoutId = setTimeout(() => {
+            // Clear the timeout reference so we can set a new one
+            timeoutId = null;
+          }, batchTimeout);
+        }
+      } else if (event.type === 'complete') {
+        // Flush remaining
+        const flushed = flushBatch();
+        if (flushed) yield flushed;
+      }
+    }
+
+    // Final flush
+    const final = flushBatch();
+    if (final) yield final;
+  }
+
+  return {
+    ...baseResult,
+    batches: createBatchStream(),
+  };
+}
+
+// -----------------------------------------------------------------------------
+// Factory Functions
+// -----------------------------------------------------------------------------
+
+/**
+ * Create an element streamer for a specific schema.
+ * Useful for reusing the same schema across multiple prompts.
+ *
+ * @example
+ * ```typescript
+ * const streamTasks = createElementStreamer({
+ *   model: anthropic('claude-sonnet'),
+ *   elementSchema: taskSchema,
+ * });
+ *
+ * // Use multiple times
+ * const webAppTasks = await streamTasks('Break down building a web app');
+ * const apiTasks = await streamTasks('Break down building an API');
+ * ```
+ */
+export function createElementStreamer<TElement extends ZodTypeAny>(config: {
+  model: SDKLanguageModelV3;
+  elementSchema: TElement;
+  system?: string;
+  maxTokens?: number;
+  temperature?: number;
+}) {
+  return async function stream(
+    prompt: string,
+    options?: Partial<ElementStreamOptions<TElement>>,
+  ): Promise<z.infer<TElement>[]> {
+    return collectElements({
+      ...config,
+      prompt,
+      ...options,
+    });
+  };
+}
+
+/**
+ * Create a collected element stream that returns all elements at once.
+ * Simpler API when you don't need real-time updates.
+ *
+ * @example
+ * ```typescript
+ * const tasks = await collectElements({
+ *   model,
+ *   elementSchema: taskSchema,
+ *   prompt: 'Generate 5 tasks',
+ * });
+ *
+ * console.log(`Got ${tasks.length} tasks`);
+ * ```
+ */
+export async function collectElements<TElement extends ZodTypeAny>(
+  options: ElementStreamOptions<TElement>,
+): Promise<z.infer<TElement>[]> {
+  const { stream, elements } = streamElements(options);
+  // Drain stream to resolve elements promise
+  for await (const _ of stream) {
+    // Consume stream to ensure elements promise resolves
+  }
+  return elements;
+}
+
+// -----------------------------------------------------------------------------
+// withElementStream Composable Wrapper
+// -----------------------------------------------------------------------------
+
+/**
+ * Configuration for withElementStream wrapper.
+ */
+export interface WithElementStreamConfig<T> {
+  /** Called for each streamed element */
+  readonly onElement?: (element: T, index: number) => void;
+  /** Buffer size before yielding (batch elements) */
+  readonly bufferSize?: number;
+  /** Transform each element before yielding */
+  readonly transform?: (element: T) => T;
+  /** Called when all elements are complete */
+  readonly onComplete?: (elements: T[], totalCount: number) => void;
+  /** Called on error */
+  readonly onError?: (error: Error) => void;
+}
+
+/**
+ * Stream result with experimental_output from streamText.
+ * This interface represents the shape of results from streamText
+ * when using experimental_output with Output.array().
+ */
+export interface ExperimentalOutputStreamResult<T> {
+  /** Element stream from experimental_output */
+  experimental_output?: {
+    /** Async iterable of elements */
+    elementStream?: AsyncIterable<T>;
+    /** Promise resolving to all elements */
+    object?: Promise<T[]>;
+  };
+}
+
+/**
+ * Composable wrapper for streaming array elements from experimental_output.
+ *
+ * This wraps a streamText result that uses experimental_output with Output.array(),
+ * providing callbacks, transformations, and buffering capabilities.
+ *
+ * @param streamResult - Result from streamText with experimental_output
+ * @param config - Configuration for element streaming
+ * @returns Async generator yielding elements (potentially transformed/buffered)
+ *
+ * @example
+ * ```typescript
+ * import { streamText, Output } from 'ai';
+ * import { withElementStream } from '@od-oneapp/ai-platform/output';
+ * import { z } from 'zod';
+ *
+ * const TodoSchema = z.object({
+ *   task: z.string(),
+ *   priority: z.enum(['low', 'medium', 'high']),
+ * });
+ *
+ * const result = streamText({
+ *   model,
+ *   prompt: 'List 10 todos',
+ *   experimental_output: Output.array({ element: TodoSchema }),
+ * });
+ *
+ * // Basic iteration with callback
+ * const todos: z.infer<typeof TodoSchema>[] = [];
+ * for await (const todo of withElementStream(result, {
+ *   onElement: (t, i) => console.log(`Got todo ${i}: ${t.task}`),
+ * })) {
+ *   todos.push(todo);
+ * }
+ *
+ * // With transformation
+ * for await (const todo of withElementStream(result, {
+ *   transform: (t) => ({ ...t, task: t.task.toUpperCase() }),
+ * })) {
+ *   console.log(todo.task); // UPPERCASE
+ * }
+ *
+ * // With buffering (yields in batches)
+ * for await (const todoBatch of withElementStream(result, {
+ *   bufferSize: 5,
+ *   onElement: (t, i) => console.log(`Buffered todo ${i}`),
+ * })) {
+ *   // todoBatch is an array of 5 todos (or fewer at the end)
+ *   processBatch(todoBatch);
+ * }
+ * ```
+ */
+export async function* withElementStream<T>(
+  streamResult: ExperimentalOutputStreamResult<T>,
+  config?: WithElementStreamConfig<T>,
+): AsyncGenerator<T | T[], void, unknown> {
+  const { onElement, bufferSize, transform, onComplete, onError } = config ?? {};
+
+  const elementStream = streamResult.experimental_output?.elementStream;
+
+  if (!elementStream) {
+    // If no element stream, try to get from object promise
+    const objectPromise = streamResult.experimental_output?.object;
+    if (objectPromise) {
+      try {
+        const elements = await objectPromise;
+        let index = 0;
+
+        if (bufferSize && bufferSize > 1) {
+          // Buffered mode
+          let buffer: T[] = [];
+          for (const element of elements) {
+            const transformed = transform ? transform(element) : element;
+            onElement?.(transformed, index);
+            buffer.push(transformed);
+            index++;
+
+            if (buffer.length >= bufferSize) {
+              yield buffer;
+              buffer = [];
+            }
+          }
+
+          // Yield remaining buffer
+          if (buffer.length > 0) {
+            yield buffer;
+          }
+        } else {
+          // Non-buffered mode
+          for (const element of elements) {
+            const transformed = transform ? transform(element) : element;
+            onElement?.(transformed, index);
+            yield transformed;
+            index++;
+          }
+        }
+
+        const finalElements = transform ? elements.map(transform) : elements;
+        onComplete?.(finalElements, finalElements.length);
+      } catch (error) {
+        const err = error instanceof Error ? error : new Error(String(error));
+        onError?.(err);
+        throw err;
+      }
+    }
+    return;
+  }
+
+  try {
+    const allElements: T[] = [];
+    let index = 0;
+    let buffer: T[] = [];
+
+    for await (const element of elementStream) {
+      const transformed = transform ? transform(element) : element;
+      allElements.push(transformed);
+
+      // Call onElement callback
+      onElement?.(transformed, index);
+      index++;
+
+      if (bufferSize && bufferSize > 1) {
+        // Buffered mode
+        buffer.push(transformed);
+
+        if (buffer.length >= bufferSize) {
+          yield buffer;
+          buffer = [];
+        }
+      } else {
+        // Non-buffered mode
+        yield transformed;
+      }
+    }
+
+    // Yield remaining buffer in buffered mode
+    if (bufferSize && bufferSize > 1 && buffer.length > 0) {
+      yield buffer;
+    }
+
+    // Call onComplete with all elements
+    onComplete?.(allElements, allElements.length);
+  } catch (error) {
+    const err = error instanceof Error ? error : new Error(String(error));
+    onError?.(err);
+    throw err;
+  }
+}
+
+/**
+ * Collect all elements from withElementStream into an array.
+ * Convenience function for when you want all elements at once.
+ *
+ * @param streamResult - Result from streamText with experimental_output
+ * @param config - Configuration for element streaming
+ * @returns Promise resolving to array of all elements
+ *
+ * @example
+ * ```typescript
+ * const todos = await collectFromElementStream(result, {
+ *   onElement: (t, i) => console.log(`Processing ${i}`),
+ *   transform: (t) => ({ ...t, processed: true }),
+ * });
+ * ```
+ */
+export async function collectFromElementStream<T>(
+  streamResult: ExperimentalOutputStreamResult<T>,
+  config?: Omit<WithElementStreamConfig<T>, 'bufferSize'>,
+): Promise<T[]> {
+  const elements: T[] = [];
+
+  for await (const element of withElementStream(streamResult, config)) {
+    // Without bufferSize, elements are yielded one at a time
+    elements.push(element as T);
+  }
+
+  return elements;
+}