feat: response source nodes in query tool output (#1784)

.changeset/two-yaks-protect.md

@@ -0,0 +1,6 @@
+---
+"llamaindex": patch
+"@llamaindex/server": patch
+---
+
+feat: response source nodes in query tool output

examples/agent/query_tool.ts

@@ -0,0 +1,51 @@
+import {
+  AgentStream,
+  AgentToolCallResult,
+  Document,
+  VectorStoreIndex,
+  agent,
+  openai,
+} from "llamaindex";
+
+async function main() {
+  const index = await VectorStoreIndex.fromDocuments([
+    new Document({
+      text: "Cats have a specialized collarbone that allows them to always land on their feet when they fall.",
+    }),
+    new Document({
+      text: "Dogs have a sense of smell that is 10,000 to 100,000 times more acute than humans.",
+    }),
+    new Document({
+      text: "Cats are known for their agility and ability to jump high.",
+    }),
+  ]);
+
+  const myAgent = agent({
+    llm: openai({ model: "gpt-4o" }),
+    tools: [
+      index.queryTool({
+        options: { similarityTopK: 2 },
+        includeSourceNodes: true,
+      }),
+    ],
+  });
+
+  const context = myAgent.run("The fact about cats");
+
+  for await (const event of context) {
+    if (event instanceof AgentToolCallResult) {
+      console.log(
+        "Using these retrieved information to answer the question:\n",
+        event.data.toolOutput.result,
+      );
+    } else if (event instanceof AgentStream) {
+      for (const chunk of event.data.delta) {
+        process.stdout.write(chunk);
+      }
+    }
+  }
+}
+
+void main().then(() => {
+  console.log("Done");
+});

packages/llamaindex/src/indices/BaseIndex.ts

@@ -41,6 +41,7 @@ export type QueryToolParams = (
 ) & {
   responseSynthesizer?: BaseSynthesizer;
   metadata?: ToolMetadata<JSONSchemaType<QueryEngineParam>> | undefined;
+  includeSourceNodes?: boolean;
 };
 
 /**
@@ -98,6 +99,7 @@ export abstract class BaseIndex<T> {
     return new QueryEngineTool({
       queryEngine: this.asQueryEngine(params),
       metadata: params?.metadata,
+      includeSourceNodes: params?.includeSourceNodes ?? false,
     });
   }
 

packages/llamaindex/src/tools/QueryEngineTool.ts

@@ -1,3 +1,4 @@
+import type { JSONValue } from "@llamaindex/core/global";
 import type { BaseTool, ToolMetadata } from "@llamaindex/core/llms";
 import type { BaseQueryEngine } from "@llamaindex/core/query-engine";
 import type { JSONSchemaType } from "ajv";
@@ -20,6 +21,7 @@ const DEFAULT_PARAMETERS: JSONSchemaType<QueryEngineParam> = {
 export type QueryEngineToolParams = {
   queryEngine: BaseQueryEngine;
   metadata?: ToolMetadata<JSONSchemaType<QueryEngineParam>> | undefined;
+  includeSourceNodes?: boolean;
 };
 
 export type QueryEngineParam = {
@@ -29,19 +31,32 @@ export type QueryEngineParam = {
 export class QueryEngineTool implements BaseTool<QueryEngineParam> {
   private queryEngine: BaseQueryEngine;
   metadata: ToolMetadata<JSONSchemaType<QueryEngineParam>>;
+  includeSourceNodes: boolean;
 
-  constructor({ queryEngine, metadata }: QueryEngineToolParams) {
+  constructor({
+    queryEngine,
+    metadata,
+    includeSourceNodes,
+  }: QueryEngineToolParams) {
     this.queryEngine = queryEngine;
     this.metadata = {
       name: metadata?.name ?? DEFAULT_NAME,
       description: metadata?.description ?? DEFAULT_DESCRIPTION,
       parameters: metadata?.parameters ?? DEFAULT_PARAMETERS,
     };
+    this.includeSourceNodes = includeSourceNodes ?? false;
   }
 
   async call({ query }: QueryEngineParam) {
     const response = await this.queryEngine.query({ query });
 
-    return response.message.content;
+    if (!this.includeSourceNodes) {
+      return { content: response.message.content };
+    }
+
+    return {
+      content: response.message.content,
+      sourceNodes: response.sourceNodes,
+    } as unknown as JSONValue;
   }
 }

packages/server/next/app/components/ui/chat/chat-starter.tsx

@@ -3,9 +3,9 @@ import { StarterQuestions } from "@llamaindex/chat-ui/widgets";
 import { getConfig } from "../lib/utils";
 
 export function ChatStarter() {
-  const { append } = useChatUI();
+  const { append, messages } = useChatUI();
   const starterQuestions = getConfig("STARTER_QUESTIONS") ?? [];
 
-  if (starterQuestions.length === 0) return null;
+  if (starterQuestions.length === 0 || messages.length > 0) return null;
   return <StarterQuestions append={append} questions={starterQuestions} />;
 }

packages/server/package.json

@@ -60,7 +60,7 @@
     "@radix-ui/react-collapsible": "^1.0.3",
     "@radix-ui/react-select": "^2.1.1",
     "@radix-ui/react-slot": "^1.0.2",
-    "@llamaindex/chat-ui": "0.3.1",
+    "@llamaindex/chat-ui": "0.3.2",
     "class-variance-authority": "^0.7.1",
     "tailwind-merge": "^2.6.0",
     "clsx": "^2.1.1",

packages/server/src/utils/suggestion.ts

@@ -0,0 +1,62 @@
+import type { StreamData } from "ai";
+import { type ChatMessage, Settings } from "llamaindex";
+
+const NEXT_QUESTION_PROMPT = `You're a helpful assistant! Your task is to suggest the next question that user might ask. 
+Here is the conversation history
+---------------------
+{conversation}
+---------------------
+Given the conversation history, please give me 3 questions that user might ask next!
+Your answer should be wrapped in three sticks which follows the following format:
+\`\`\`
+<question 1>
+<question 2>
+<question 3>
+\`\`\`
+`;
+
+export const sendSuggestedQuestionsEvent = async (
+  dataStream: StreamData,
+  chatHistory: ChatMessage[] = [],
+) => {
+  const questions = await generateNextQuestions(chatHistory);
+  if (questions.length > 0) {
+    dataStream.appendMessageAnnotation({
+      type: "suggested_questions",
+      data: questions,
+    });
+  }
+};
+
+async function generateNextQuestions(conversation: ChatMessage[]) {
+  const conversationText = conversation
+    .map((message) => `${message.role}: ${message.content}`)
+    .join("\n");
+  const message = NEXT_QUESTION_PROMPT.replace(
+    "{conversation}",
+    conversationText,
+  );
+
+  try {
+    const response = await Settings.llm.complete({ prompt: message });
+    const questions = extractQuestions(response.text);
+    return questions;
+  } catch (error) {
+    console.error("Error when generating the next questions: ", error);
+    return [];
+  }
+}
+
+function extractQuestions(text: string): string[] {
+  // Extract the text inside the triple backticks
+  const contentMatch = text.match(/```(.*?)```/s);
+  const content = contentMatch?.[1] ?? "";
+
+  // Split the content by newlines to get each question
+  const questions = content
+    .split("\n")
+    .map((question) => question.trim())
+    .filter((question) => question !== "");
+
+  return questions;
+}

packages/server/src/utils/workflow.ts

@@ -13,6 +13,7 @@ import {
 } from "llamaindex";
 import { ReadableStream } from "stream/web";
 import type { ServerWorkflow } from "../types";
+import { sendSuggestedQuestionsEvent } from "./suggestion";
 
 export async function runWorkflow(
   workflow: ServerWorkflow,
@@ -47,11 +48,19 @@ async function runAgentWorkflow(
         }
       }
       controller.close();
-      dataStream.close();
     },
   });
 
-  return LlamaIndexAdapter.toDataStreamResponse(stream, { data: dataStream });
+  return LlamaIndexAdapter.toDataStreamResponse(stream, {
+    data: dataStream,
+    callbacks: {
+      onFinal: async (content: string) => {
+        const history = chatHistory.concat({ role: "assistant", content });
+        await sendSuggestedQuestionsEvent(dataStream, history);
+        dataStream.close();
+      },
+    },
+  });
 }
 
 async function runCustomWorkflow(
@@ -77,9 +86,20 @@ async function runCustomWorkflow(
         }
       }
       controller.close();
-      dataStream.close();
     },
   });
 
-  return LlamaIndexAdapter.toDataStreamResponse(stream, { data: dataStream });
+  return LlamaIndexAdapter.toDataStreamResponse(stream, {
+    data: dataStream,
+    callbacks: {
+      onFinal: async (content: string) => {
+        const history = agentInput.chatHistory?.concat({
+          role: "assistant",
+          content,
+        });
+        await sendSuggestedQuestionsEvent(dataStream, history);
+        dataStream.close();
+      },
+    },
+  });
 }

pnpm-lock.yaml

@@ -1767,8 +1767,8 @@ importers:
   packages/server:
     dependencies:
       '@llamaindex/chat-ui':
-        specifier: 0.3.1
-        version: 0.3.1(@types/react-dom@19.0.4(@types/react@19.0.10))(@types/react@19.0.10)(react-dom@19.0.0(react@19.0.0))(react@19.0.0)
+        specifier: 0.3.2
+        version: 0.3.2(@types/react-dom@19.0.4(@types/react@19.0.10))(@types/react@19.0.10)(react-dom@19.0.0(react@19.0.0))(react@19.0.0)
       '@radix-ui/react-accordion':
         specifier: ^1.2.2
         version: 1.2.3(@types/react-dom@19.0.4(@types/react@19.0.10))(@types/react@19.0.10)(react-dom@19.0.0(react@19.0.0))(react@19.0.0)
@@ -2445,10 +2445,6 @@ packages:
     peerDependencies:
       '@babel/core': ^7.0.0-0
 
-  '@babel/runtime@7.26.7':
-    resolution: {integrity: sha512-AOPI3D+a8dXnja+iwsUqGRjr1BbZIe771sXdapOtYI531gSqpi92vXivKcq2asu/DFpdl1ceFAKZyRzK2PCVcQ==}
-    engines: {node: '>=6.9.0'}
-
   '@babel/runtime@7.26.9':
     resolution: {integrity: sha512-aA63XwOkcl4xxQa3HjPMqOP6LiK0ZDv3mUPYEFXkpHbaFjtGggE1A61FjFzJnB+p7/oy2gA8E+rcBNl/zC1tMg==}
     engines: {node: '>=6.9.0'}
@@ -3738,8 +3734,8 @@ packages:
     peerDependencies:
       react: ^18.2.0 || ^19.0.0 || ^19.0.0-rc
 
-  '@llamaindex/chat-ui@0.3.1':
-    resolution: {integrity: sha512-sF6axN9LviewAxvBbqkF3u3K0yvIt74prio7uiVruFVT/AYkRlIk721QXTPBscf+ZvyzAqjh0Nx0BoGiZUzBCw==}
+  '@llamaindex/chat-ui@0.3.2':
+    resolution: {integrity: sha512-YcQOghcxutqHK9KO2CRSws0inDR5bbMZkmpUFJtC2aWcHjWi8wYbzVZjRVl1vrb3VCk+VInKOhFUTW9hEkzydA==}
     peerDependencies:
       react: ^18.2.0 || ^19.0.0 || ^19.0.0-rc
 
@@ -13700,10 +13696,6 @@ snapshots:
       '@babel/core': 7.26.8
       '@babel/helper-plugin-utils': 7.26.5
 
-  '@babel/runtime@7.26.7':
-    dependencies:
-      regenerator-runtime: 0.14.1
-
   '@babel/runtime@7.26.9':
     dependencies:
       regenerator-runtime: 0.14.1
@@ -14958,7 +14950,7 @@ snapshots:
       - react-dom
       - supports-color
 
-  '@llamaindex/chat-ui@0.3.1(@types/react-dom@19.0.4(@types/react@19.0.10))(@types/react@19.0.10)(react-dom@19.0.0(react@19.0.0))(react@19.0.0)':
+  '@llamaindex/chat-ui@0.3.2(@types/react-dom@19.0.4(@types/react@19.0.10))(@types/react@19.0.10)(react-dom@19.0.0(react@19.0.0))(react@19.0.0)':
     dependencies:
       '@llamaindex/pdf-viewer': 1.3.0(@types/react@19.0.10)(react-dom@19.0.0(react@19.0.0))(react@19.0.0)
       '@radix-ui/react-collapsible': 1.1.3(@types/react-dom@19.0.4(@types/react@19.0.10))(@types/react@19.0.10)(react-dom@19.0.0(react@19.0.0))(react@19.0.0)
@@ -15005,14 +14997,14 @@ snapshots:
 
   '@manypkg/find-root@1.1.0':
     dependencies:
-      '@babel/runtime': 7.26.7
+      '@babel/runtime': 7.26.9
       '@types/node': 12.20.55
       find-up: 4.1.0
       fs-extra: 8.1.0
 
   '@manypkg/get-packages@1.1.3':
     dependencies:
-      '@babel/runtime': 7.26.7
+      '@babel/runtime': 7.26.9
       '@changesets/types': 4.1.0
       '@manypkg/find-root': 1.1.0
       fs-extra: 8.1.0