Add adapter to use Vercel model providers (#1558)

Co-authored-by: Alex Yang <himself65@outlook.com>

.changeset/hip-windows-tease.md

@@ -0,0 +1,5 @@
+---
+"@llamaindex/vercel": patch
+---
+
+Add VercelLLM (adapter to use any model provider from Vercel AI in LlamaIndex)

apps/next/src/content/docs/llamaindex/integration/vercel.mdx

@@ -3,7 +3,9 @@ title: Vercel
 description: Integrate LlamaIndex with Vercel's AI SDK
 ---
 
-LlamaIndex provides integration with Vercel's AI SDK, allowing you to create powerful search and retrieval applications. Below are examples of how to use LlamaIndex with `streamText` from the Vercel AI SDK.
+LlamaIndex provides integration with Vercel's AI SDK, allowing you to create powerful search and retrieval applications. You can:
+- Use any of Vercel AI's [model providers](https://sdk.vercel.ai/docs/foundations/providers-and-models) as LLMs in LlamaIndex
+- Use indexes (e.g. VectorStoreIndex, LlamaCloudIndex) from LlamaIndexTS in your Vercel AI applications
 
 ## Setup
 
@@ -13,7 +15,22 @@ First, install the required dependencies:
 npm install @llamaindex/vercel ai
 ```
 
-## Using Local Vector Store
+## Using Vercel AI's Model Providers 
+
+Using the `VercelLLM` adapter, it's easy to use any of Vercel AI's [model providers](https://sdk.vercel.ai/docs/foundations/providers-and-models) as LLMs in LlamaIndex. Here's an example of how to use OpenAI's GPT-4o model:
+
+```typescript
+const llm = new VercelLLM({ model: openai("gpt-4o") });
+const result = await llm.complete({
+  prompt: "What is the capital of France?",
+  stream: false, // Set to true if you want streaming responses
+});
+console.log(result.text);
+```
+
+## Use Indexes
+
+### Using VectorStoreIndex
 
 Here's how to create a simple vector store index and query it using Vercel's AI SDK:
 
@@ -29,22 +46,25 @@ const index = await VectorStoreIndex.fromDocuments([document]);
 
 // Create a query tool
 const queryTool = llamaindex({
+  model: openai("gpt-4"),
   index,
   description: "Search through the documents", // optional
 });
 
 // Use the tool with Vercel's AI SDK
 streamText({
-  tools: { queryTool },
-  prompt: "Your question here",
   model: openai("gpt-4"),
+  prompt: "Your question here",
+  tools: { queryTool },
   onFinish({ response }) {
     console.log("Response:", response.messages); // log the response
   },
 }).toDataStream();
 ```
 
-## Using LlamaCloud
+> Note: the Vercel AI model referenced in the `llamaindex` function is used by the response synthesizer to generate a response for the tool call.
+
+### Using LlamaCloud
 
 For production deployments, you can use LlamaCloud to store and manage your documents:
 
@@ -61,15 +81,16 @@ const index = await LlamaCloudIndex.fromDocuments({
 
 // Use it the same way as VectorStoreIndex
 const queryTool = llamaindex({
+  model: openai("gpt-4"),
   index,
   description: "Search through the documents",
 });
 
 // Use the tool with Vercel's AI SDK
 streamText({
-  tools: { queryTool },
-  prompt: "Your question here",
   model: openai("gpt-4"),
+  prompt: "Your question here",
+  tools: { queryTool },
 }).toDataStream();
 ```
 

examples/vercel/README.md

@@ -14,6 +14,16 @@ npm i
 
 Make sure to run the examples from the parent folder called `examples`. The following examples are available:
 
+### Vercel LLM Example
+
+Run the Vercel LLM example with:
+
+```bash
+npx tsx vercel/llm.ts
+```
+
+This example demonstrates using the `VercelLLM` adapter with Vercel's OpenAI model provider
+
 ### Vector Store Example
 
 Run the local vector store example with:

examples/vercel/llamacloud.ts

@@ -22,6 +22,7 @@ async function main() {
     prompt: "Cost of moving cat from Russia to UK?",
     tools: {
       queryTool: llamaindex({
+        model: openai("gpt-4o"),
         index,
         description:
           "get information from your knowledge base to answer questions.", // optional description

examples/vercel/llm.ts

@@ -0,0 +1,45 @@
+import { openai } from "@ai-sdk/openai";
+import { VercelLLM } from "@llamaindex/vercel";
+import { LLMAgent, WikipediaTool } from "llamaindex";
+
+async function main() {
+  // Create an instance of VercelLLM with the OpenAI model
+  const vercelLLM = new VercelLLM({ model: openai("gpt-4o") });
+
+  console.log("\n=== Test 1: Using complete() for single response ===");
+  const result = await vercelLLM.complete({
+    prompt: "What is the capital of France?",
+    stream: false, // Set to true if you want streaming responses
+  });
+  console.log(result.text);
+
+  console.log("\n=== Test 2: Using chat() for streaming response ===");
+  const stream = await vercelLLM.chat({
+    messages: [
+      { content: "You want to talk in rhymes.", role: "system" },
+      {
+        content:
+          "How much wood would a woodchuck chuck if a woodchuck could chuck wood?",
+        role: "user",
+      },
+    ],
+    stream: true,
+  });
+  for await (const chunk of stream) {
+    process.stdout.write(chunk.delta);
+  }
+
+  console.log("\n=== Test 3: Using LLMAgent with WikipediaTool ===");
+  const agent = new LLMAgent({
+    llm: vercelLLM,
+    tools: [new WikipediaTool()],
+  });
+
+  const { message } = await agent.chat({
+    message: "What's the history of New York from Wikipedia in 3 sentences?",
+  });
+
+  console.log(message);
+}
+
+main().catch(console.error);

examples/vercel/vector-store.ts

@@ -18,6 +18,7 @@ async function main() {
     prompt: "Cost of moving cat from Russia to UK?",
     tools: {
       queryTool: llamaindex({
+        model: openai("gpt-4o"),
         index,
         description:
           "get information from your knowledge base to answer questions.", // optional description

packages/providers/vercel/src/index.ts

@@ -1 +1,2 @@
+export { VercelLLM } from "./llm";
 export { llamaindex } from "./tool";

packages/providers/vercel/src/llm.ts

@@ -0,0 +1,183 @@
+import { wrapEventCaller, wrapLLMEvent } from "@llamaindex/core/decorator";
+import {
+  ToolCallLLM,
+  type ChatMessage,
+  type ChatResponse,
+  type ChatResponseChunk,
+  type LLMChatParamsNonStreaming,
+  type LLMChatParamsStreaming,
+  type LLMMetadata,
+  type ToolCallLLMMessageOptions,
+} from "@llamaindex/core/llms";
+import { extractText } from "@llamaindex/core/utils";
+import {
+  generateText,
+  streamText,
+  type CoreAssistantMessage,
+  type CoreMessage,
+  type CoreSystemMessage,
+  type CoreToolMessage,
+  type CoreUserMessage,
+  type ImagePart,
+  type LanguageModelV1,
+  type TextPart,
+} from "ai";
+
+export type VercelAdditionalChatOptions = ToolCallLLMMessageOptions;
+
+export class VercelLLM extends ToolCallLLM<VercelAdditionalChatOptions> {
+  supportToolCall: boolean = true;
+  private model: LanguageModelV1;
+
+  constructor({ model }: { model: LanguageModelV1 }) {
+    super();
+    this.model = model;
+  }
+
+  get metadata(): LLMMetadata {
+    return {
+      model: this.model.modelId,
+      temperature: 1,
+      topP: 1,
+      contextWindow: 128000,
+      tokenizer: undefined,
+    };
+  }
+
+  private toVercelMessages(
+    messages: ChatMessage<ToolCallLLMMessageOptions>[],
+  ): CoreMessage[] {
+    return messages.map((message) => {
+      const options = message.options ?? {};
+
+      if ("toolResult" in options) {
+        return {
+          role: "tool",
+          content: [
+            {
+              type: "tool-result",
+              toolCallId: options.toolResult.id,
+              toolName: "", // XXX: tool result doesn't name
+              isError: options.toolResult.isError,
+              result: options.toolResult.result,
+            },
+          ],
+        } satisfies CoreToolMessage;
+      } else if ("toolCall" in options) {
+        return {
+          role: "assistant",
+          content: options.toolCall.map((toolCall) => ({
+            type: "tool-call",
+            toolName: toolCall.name,
+            toolCallId: toolCall.id,
+            args: toolCall.input,
+          })),
+        } satisfies CoreAssistantMessage;
+      }
+
+      if (message.role === "system" || message.role === "assistant") {
+        return {
+          role: message.role,
+          content: extractText(message.content),
+        } satisfies CoreSystemMessage | CoreAssistantMessage;
+      }
+
+      if (message.role === "user") {
+        return {
+          role: message.role,
+          content:
+            typeof message.content === "string"
+              ? message.content
+              : message.content.map((contentDetail) => {
+                  if (contentDetail.type === "image_url") {
+                    return {
+                      type: "image",
+                      image: new URL(contentDetail.image_url.url),
+                    } satisfies ImagePart;
+                  }
+                  return {
+                    type: "text",
+                    text: contentDetail.text,
+                  } satisfies TextPart;
+                }),
+        } satisfies CoreUserMessage;
+      }
+
+      throw new Error(`Can not convert message ${JSON.stringify(message)}`);
+    });
+  }
+
+  chat(
+    params: LLMChatParamsStreaming<
+      VercelAdditionalChatOptions,
+      ToolCallLLMMessageOptions
+    >,
+  ): Promise<AsyncIterable<ChatResponseChunk<ToolCallLLMMessageOptions>>>;
+  chat(
+    params: LLMChatParamsNonStreaming<
+      VercelAdditionalChatOptions,
+      ToolCallLLMMessageOptions
+    >,
+  ): Promise<ChatResponse<ToolCallLLMMessageOptions>>;
+  @wrapEventCaller
+  @wrapLLMEvent
+  async chat(
+    params:
+      | LLMChatParamsNonStreaming<
+          VercelAdditionalChatOptions,
+          ToolCallLLMMessageOptions
+        >
+      | LLMChatParamsStreaming<
+          VercelAdditionalChatOptions,
+          ToolCallLLMMessageOptions
+        >,
+  ): Promise<
+    | ChatResponse<ToolCallLLMMessageOptions>
+    | AsyncIterable<ChatResponseChunk<ToolCallLLMMessageOptions>>
+  > {
+    const { messages, stream } = params;
+
+    // Streaming
+    if (stream) {
+      const result = streamText({
+        model: this.model,
+        messages: this.toVercelMessages(messages),
+      });
+      return result.fullStream.pipeThrough(
+        new TransformStream({
+          async transform(message, controller): Promise<void> {
+            switch (message.type) {
+              case "text-delta":
+                controller.enqueue({ raw: message, delta: message.textDelta });
+            }
+          },
+        }),
+      );
+    }
+
+    // Non-streaming
+    const result = await generateText({
+      model: this.model,
+      messages: this.toVercelMessages(messages),
+    });
+
+    return {
+      raw: result,
+      message: {
+        content: result.text,
+        role: "assistant",
+        options: result.toolCalls?.length
+          ? {
+              toolCall: result.toolCalls.map(
+                ({ toolCallId, toolName, args }) => ({
+                  id: toolCallId,
+                  name: toolName,
+                  input: args,
+                }),
+              ),
+            }
+          : {},
+      },
+    };
+  }
+}

packages/providers/vercel/src/tool.ts

@@ -1,29 +1,36 @@
+import { Settings } from "@llamaindex/core/global";
 import type { BaseQueryEngine } from "@llamaindex/core/query-engine";
-import { type CoreTool, tool } from "ai";
+import { type CoreTool, type LanguageModelV1, tool } from "ai";
 import { z } from "zod";
+import { VercelLLM } from "./llm";
 
 interface DatasourceIndex {
   asQueryEngine: () => BaseQueryEngine;
 }
 
 export function llamaindex({
+  model,
   index,
   description,
 }: {
+  model: LanguageModelV1;
   index: DatasourceIndex;
   description?: string;
 }): CoreTool {
-  const queryEngine = index.asQueryEngine();
-  return tool({
-    description: description ?? "Get information about your documents.",
-    parameters: z.object({
-      query: z
-        .string()
-        .describe("The query to get information about your documents."),
-    }),
-    execute: async ({ query }) => {
-      const result = await queryEngine?.query({ query });
-      return result?.message.content ?? "No result found in documents.";
-    },
+  const llm = new VercelLLM({ model });
+  return Settings.withLLM<CoreTool>(llm, () => {
+    const queryEngine = index.asQueryEngine();
+    return tool({
+      description: description ?? "Get information about your documents.",
+      parameters: z.object({
+        query: z
+          .string()
+          .describe("The query to get information about your documents."),
+      }),
+      execute: async ({ query }) => {
+        const result = await queryEngine?.query({ query });
+        return result?.message.content ?? "No result found in documents.";
+      },
+    });
   });
 }