feat: vercel tool response fields options (#1765)

.changeset/purple-gifts-decide.md

@@ -0,0 +1,6 @@
+---
+"@llamaindex/vercel": minor
+"@llamaindex/doc": minor
+---
+
+Adding an options parameter to vercel tool to tailor responses

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

@@ -84,6 +84,7 @@ const queryTool = llamaindex({
   model: openai("gpt-4"),
   index,
   description: "Search through the documents",
+  options: { fields: ["sourceNodes", "messages"]}
 });
 
 // Use the tool with Vercel's AI SDK

packages/providers/vercel/package.json

@@ -35,10 +35,12 @@
   },
   "scripts": {
     "build": "bunchee",
-    "dev": "bunchee --watch"
+    "dev": "bunchee --watch",
+    "test": "vitest run"
   },
   "devDependencies": {
-    "bunchee": "6.4.0"
+    "bunchee": "6.4.0",
+    "vitest": "^2.1.5"
   },
   "dependencies": {
     "@llamaindex/core": "workspace:*",

packages/providers/vercel/src/tool.ts

@@ -1,5 +1,6 @@
 import { Settings } from "@llamaindex/core/global";
 import type { BaseQueryEngine } from "@llamaindex/core/query-engine";
+import { EngineResponse } from "@llamaindex/core/schema";
 import { type CoreTool, type LanguageModelV1, tool } from "ai";
 import { z } from "zod";
 import { VercelLLM } from "./llm";
@@ -8,14 +9,20 @@ interface DatasourceIndex {
   asQueryEngine: () => BaseQueryEngine;
 }
 
+type ResponseField = keyof EngineResponse;
+
 export function llamaindex({
   model,
   index,
   description,
+  options,
 }: {
   model: LanguageModelV1;
   index: DatasourceIndex;
   description?: string;
+  options?: {
+    fields?: ResponseField[];
+  };
 }): CoreTool {
   const llm = new VercelLLM({ model });
   return Settings.withLLM<CoreTool>(llm, () => {
@@ -29,6 +36,15 @@ export function llamaindex({
       }),
       execute: async ({ query }) => {
         const result = await queryEngine?.query({ query });
+        if (options?.fields) {
+          const resultWithFields = {
+            ...result,
+            ...Object.fromEntries(
+              options.fields.map((field) => [field, result[field]]),
+            ),
+          };
+          return resultWithFields;
+        }
         return result?.message.content ?? "No result found in documents.";
       },
     });

packages/providers/vercel/tests/index.test.ts

@@ -0,0 +1,168 @@
+import {
+  type CoreTool,
+  type LanguageModelV1,
+  type ToolExecutionOptions,
+} from "ai";
+import { beforeAll, describe, expect, test, vi } from "vitest";
+import { llamaindex } from "../src/tool";
+
+// Mock the ai package functions
+vi.mock("ai", () => ({
+  generateText: vi.fn().mockResolvedValue({
+    text: "Hi there!",
+    reasoning: "",
+    sources: [],
+    experimental_output: {},
+    toolCalls: [],
+    messages: [],
+    data: {},
+    usage: {},
+    id: "test",
+    createdAt: new Date(),
+  }),
+  streamText: vi.fn(),
+  tool: vi.fn().mockImplementation((config) => ({
+    name: config.name || "llamaindex",
+    description: config.description,
+    execute: config.execute,
+  })),
+}));
+
+describe("llamaindex Tool", () => {
+  let mockModel: LanguageModelV1;
+  let mockToolOptions: ToolExecutionOptions;
+
+  beforeAll(() => {
+    mockModel = {
+      modelId: "test-model",
+    } as LanguageModelV1;
+
+    mockToolOptions = {
+      toolCallId: "test-call-id",
+      messages: [],
+    };
+  });
+
+  test("creates a tool with default description", () => {
+    const mockIndex = {
+      asQueryEngine: vi.fn().mockReturnValue({
+        query: vi.fn().mockResolvedValue({
+          message: { content: "Test response" },
+        }),
+      }),
+    };
+
+    const tool = llamaindex({
+      model: mockModel,
+      index: mockIndex,
+    });
+
+    expect(tool).toBeDefined();
+    expect(typeof tool.execute).toBe("function");
+  });
+
+  test("creates a tool with custom description", () => {
+    const mockIndex = {
+      asQueryEngine: vi.fn().mockReturnValue({
+        query: vi.fn().mockResolvedValue({
+          message: { content: "Test response" },
+        }),
+      }),
+    };
+
+    const tool = llamaindex({
+      model: mockModel,
+      index: mockIndex,
+      description: "Custom description",
+    });
+
+    expect(tool).toBeDefined();
+    expect(typeof tool.execute).toBe("function");
+  });
+
+  describe("Tool Execution", () => {
+    test("execute returns message content when no options specified", async () => {
+      const mockQueryResult = {
+        message: { content: "Test response" },
+        sourceNodes: [{ text: "Source 1" }],
+        metadata: { some: "data" },
+      };
+
+      const mockIndex = {
+        asQueryEngine: vi.fn().mockReturnValue({
+          query: vi.fn().mockResolvedValue(mockQueryResult),
+        }),
+      };
+
+      const tool = llamaindex({
+        model: mockModel,
+        index: mockIndex,
+      });
+
+      // Ensure tool.execute exists before calling it
+      expect(tool.execute).toBeDefined();
+      const result = await (tool as Required<CoreTool>).execute(
+        { query: "test query" },
+        mockToolOptions,
+      );
+      expect(result).toBe("Test response");
+    });
+
+    test("execute returns specified fields when options provided", async () => {
+      const mockQueryResult = {
+        message: { content: "Test response" },
+        sourceNodes: [{ text: "Source 1" }],
+        metadata: { some: "data" },
+      };
+
+      const mockIndex = {
+        asQueryEngine: vi.fn().mockReturnValue({
+          query: vi.fn().mockResolvedValue(mockQueryResult),
+        }),
+      };
+
+      const tool = llamaindex({
+        model: mockModel,
+        index: mockIndex,
+        options: {
+          fields: ["sourceNodes", "metadata"],
+        },
+      });
+
+      // Ensure tool.execute exists before calling it
+      expect(tool.execute).toBeDefined();
+      const result = await (tool as Required<CoreTool>).execute(
+        { query: "test query" },
+        mockToolOptions,
+      );
+      expect(result).toEqual({
+        message: { content: "Test response" },
+        sourceNodes: [{ text: "Source 1" }],
+        metadata: { some: "data" },
+      });
+    });
+
+    test("execute returns 'No result found' when query returns no content", async () => {
+      const mockIndex = {
+        asQueryEngine: vi.fn().mockReturnValue({
+          query: vi.fn().mockResolvedValue({
+            message: { content: null },
+          }),
+        }),
+      };
+
+      const tool = llamaindex({
+        model: mockModel,
+        index: mockIndex,
+      });
+
+      // Ensure tool.execute exists before calling it
+      expect(tool.execute).toBeDefined();
+      const result = await (tool as Required<CoreTool>).execute(
+        { query: "test query" },
+        mockToolOptions,
+      );
+      expect(result).toBe("No result found in documents.");
+    });
+  });
+});

pnpm-lock.yaml

@@ -1682,6 +1682,9 @@ importers:
       bunchee:
         specifier: 6.4.0
         version: 6.4.0(typescript@5.7.3)
+      vitest:
+        specifier: ^2.1.5
+        version: 2.1.5(@edge-runtime/vm@4.0.4)(@types/node@22.13.5)(happy-dom@15.11.7)(lightningcss@1.29.1)(msw@2.7.0(@types/node@22.13.5)(typescript@5.7.3))(terser@5.38.2)
 
   packages/providers/vllm:
     dependencies: