Add support for LlamaIndex

.changeset/silent-days-lay.md

@@ -0,0 +1,5 @@
+---
+'ai': patch
+---
+
+Add support for LlamaIndex

content/docs/07-reference/stream-helpers/16-llamaindex-adapter.mdx

@@ -0,0 +1,61 @@
+---
+title: LlamaIndexAdapter
+description: API Reference for LlamaIndexAdapter.
+---
+
+# `LlamaIndexAdapter`
+
+The `LlamaIndexAdapter` module provides helper functions to transform LlamaIndex output streams into data streams and data stream responses.
+See the [LlamaIndex Adapter documentation](/providers/adapters/llamaindex) for more information.
+
+It supports:
+
+- LlamaIndex ChatEngine streams
+- LlamaIndex QueryEngine streams
+
+## Import
+
+<Snippet text={`import { LlamaIndexAdapter } from "ai"`} prompt={false} />
+
+## API Signature
+
+### Methods
+
+<PropertiesTable
+  content={[
+    {
+      name: 'toDataStream',
+      type: '(stream: AsyncIterable<EngineResponse>, AIStreamCallbacksAndOptions) => AIStream',
+      description: 'Converts LlamaIndex output streams to data stream.',
+    },
+    {
+      name: 'toDataStreamResponse',
+      type: '(stream: AsyncIterable<EngineResponse>, options?: {init?: ResponseInit, data?: StreamData, callbacks?: AIStreamCallbacksAndOptions}) => Response',
+      description:
+        'Converts LlamaIndex output streams to data stream response.',
+    },
+  ]}
+/>
+
+## Examples
+
+### Convert LlamaIndex ChatEngine Stream
+
+```tsx filename="app/api/completion/route.ts" highlight={"14"}
+import { OpenAI, SimpleChatEngine } from 'llamaindex';
+import { LlamaIndexAdapter } from 'ai';
+
+export async function POST(req: Request) {
+  const { prompt } = await req.json();
+
+  const llm = new OpenAI({ model: 'gpt-4o' });
+  const chatEngine = new SimpleChatEngine({ llm });
+
+  const stream = await chatEngine.chat({
+    message: prompt,
+    stream: true,
+  });
+
+  return LlamaIndexAdapter.toDataStreamResponse(stream);
+}
+```

content/docs/07-reference/stream-helpers/index.mdx

@@ -91,6 +91,12 @@ description: Learn to use help functions that help stream generations from diffe
       description:
         "Transforms the response from LangChain's stream into AI streams.",
       href: '/docs/reference/stream-helpers/langchain-adapter',
+    },
+     {
+      title: 'LlamaIndexAdapter',
+      description:
+        "Transforms the response from LlamaIndex's streams into AI streams.",
+      href: '/docs/reference/stream-helpers/llamaindex-adapter',
     },
     {
       title: 'MistralStream',

content/providers/04-adapters/index.mdx

@@ -12,3 +12,4 @@ with 3rd party libraries.
 The following adapters are currently available:
 
 - [LangChain](/providers/adapters/langchain)
+- [LlamaIndex](/providers/adapters/llamaindex)

content/providers/04-adapters/llamaindex.mdx

@@ -0,0 +1,61 @@
+---
+title: LlamaIndex
+description: Learn how to use LlamaIndex with the Vercel AI SDK.
+---
+
+# LlamaIndex
+
+[LlamaIndex](https://ts.llamaindex.ai/) is a framework for building LLM-powered applications. LlamaIndex helps you ingest, structure, and access private or domain-specific data. LlamaIndex.TS offers the core features of LlamaIndex for Python for popular runtimes like Node.js (official support), Vercel Edge Functions (experimental), and Deno (experimental).
+
+## Example: Completion
+
+Here is a basic example that uses both Vercel AI SDK and LlamaIndex together with the [Next.js](https://nextjs.org/docs) App Router.
+
+The AI SDK [`LlamaIndexAdapter`](/docs/reference/stream-helpers/llamaindex-adapter) uses the stream result from calling the `chat` method on a [LlamaIndex ChatEngine](https://ts.llamaindex.ai/modules/chat_engine) or the `query` method on a [LlamaIndex QueryEngine](https://ts.llamaindex.ai/modules/query_engines) to pipe text to the client.
+
+```tsx filename="app/api/completion/route.ts" highlight={"16"}
+import { OpenAI, SimpleChatEngine } from 'llamaindex';
+import { LlamaIndexAdapter } from 'ai';
+
+export const maxDuration = 60;
+
+export async function POST(req: Request) {
+  const { prompt } = await req.json();
+
+  const llm = new OpenAI({ model: 'gpt-4o' });
+  const chatEngine = new SimpleChatEngine({ llm });
+
+  const stream = await chatEngine.chat({
+    message: prompt,
+    stream: true,
+  });
+
+  return LlamaIndexAdapter.toDataStreamResponse(stream);
+}
+```
+
+Then, we use the Vercel AI SDK's [`useCompletion`](/docs/ai-sdk-ui/completion) method in the page component to handle the completion:
+
+```tsx filename="app/page.tsx"
+'use client';
+
+import { useCompletion } from 'ai/react';
+
+export default function Chat() {
+  const { completion, input, handleInputChange, handleSubmit } =
+    useCompletion();
+
+  return (
+    <div>
+      {completion}
+      <form onSubmit={handleSubmit}>
+        <input value={input} onChange={handleInputChange} />
+      </form>
+    </div>
+  );
+}
+```
+
+## More Examples
+
+[create-llama](https://github.com/run-llama/create-llama) is the easiest way to get started with LlamaIndex. It uses the Vercel AI SDK to connect to LlamaIndex in all its generated code.

packages/ai/streams/index.ts

@@ -50,6 +50,7 @@ export * from './google-generative-ai-stream';
 export * from './huggingface-stream';
 export * from './inkeep-stream';
 export * as LangChainAdapter from './langchain-adapter';
+export * as LlamaIndexAdapter from './llamaindex-adapter';
 export * from './langchain-stream';
 export * from './mistral-stream';
 export * from './openai-stream';

packages/ai/streams/llamaindex-adapter.test.ts

@@ -0,0 +1,50 @@
+import {
+  convertReadableStreamToArray,
+  convertResponseStreamToArray,
+  convertArrayToAsyncIterable,
+} from '@ai-sdk/provider-utils/test';
+import { toDataStream, toDataStreamResponse } from './llamaindex-adapter';
+
+describe('toDataStream', () => {
+  it('should convert AsyncIterable<EngineResponse>', async () => {
+    const inputStream = convertArrayToAsyncIterable([
+      { delta: 'Hello' },
+      { delta: 'World' },
+    ]);
+
+    assert.deepStrictEqual(
+      await convertReadableStreamToArray(
+        toDataStream(inputStream).pipeThrough(new TextDecoderStream()),
+      ),
+      ['0:"Hello"\n', '0:"World"\n'],
+    );
+  });
+});
+
+describe('toDataStreamResponse', () => {
+  it('should convert AsyncIterable<EngineResponse>', async () => {
+    const inputStream = convertArrayToAsyncIterable([
+      { delta: 'Hello' },
+      { delta: 'World' },
+    ]);
+
+    const response = toDataStreamResponse(inputStream);
+
+    assert.strictEqual(response.status, 200);
+
+    assert.deepStrictEqual(Object.fromEntries(response.headers.entries()), {
+      'content-type': 'text/plain; charset=utf-8',
+      'x-vercel-ai-data-stream': 'v1',
+    });
+
+    assert.strictEqual(
+      response.headers.get('Content-Type'),
+      'text/plain; charset=utf-8',
+    );
+
+    assert.deepStrictEqual(await convertResponseStreamToArray(response), [
+      '0:"Hello"\n',
+      '0:"World"\n',
+    ]);
+  });
+});

packages/ai/streams/llamaindex-adapter.ts

@@ -0,0 +1,64 @@
+import { mergeStreams } from '../core/util/merge-streams';
+import { prepareResponseHeaders } from '../core/util/prepare-response-headers';
+import { createStreamDataTransformer, StreamData } from './stream-data';
+import {
+  AIStreamCallbacksAndOptions,
+  createCallbacksTransformer,
+  trimStartOfStreamHelper,
+} from './ai-stream';
+
+type EngineResponse = {
+  delta: string;
+};
+
+export function toDataStream(
+  stream: AsyncIterable<EngineResponse>,
+  callbacks?: AIStreamCallbacksAndOptions,
+) {
+  return toReadableStream(stream)
+    .pipeThrough(createCallbacksTransformer(callbacks))
+    .pipeThrough(createStreamDataTransformer());
+}
+
+export function toDataStreamResponse(
+  stream: AsyncIterable<EngineResponse>,
+  options: {
+    init?: ResponseInit;
+    data?: StreamData;
+    callbacks?: AIStreamCallbacksAndOptions;
+  } = {},
+) {
+  const { init, data, callbacks } = options;
+  const dataStream = toDataStream(stream, callbacks);
+  const responseStream = data
+    ? mergeStreams(data.stream, dataStream)
+    : dataStream;
+
+  return new Response(responseStream, {
+    status: init?.status ?? 200,
+    statusText: init?.statusText,
+    headers: prepareResponseHeaders(init, {
+      contentType: 'text/plain; charset=utf-8',
+      dataStreamVersion: 'v1',
+    }),
+  });
+}
+
+function toReadableStream(res: AsyncIterable<EngineResponse>) {
+  const it = res[Symbol.asyncIterator]();
+  const trimStartOfStream = trimStartOfStreamHelper();
+
+  return new ReadableStream<string>({
+    async pull(controller): Promise<void> {
+      const { value, done } = await it.next();
+      if (done) {
+        controller.close();
+        return;
+      }
+      const text = trimStartOfStream(value.delta ?? '');
+      if (text) {
+        controller.enqueue(text);
+      }
+    },
+  });
+}