refactor: unify response and agent response

packages/core/src/Response.ts

@@ -1,16 +1,27 @@
 import type { NodeWithScore } from "./Node.js";
+import type { ChatMessage, ChatResponse } from "./llm/types.js";
 
-/**
- * Response is the output of a LLM
- */
-export class Response {
+export class Response implements ChatResponse {
+  // @deprecated use 'message' instead
   response: string;
   sourceNodes?: NodeWithScore[];
   metadata: Record<string, unknown> = {};
 
-  constructor(response: string, sourceNodes?: NodeWithScore[]) {
+  message: ChatMessage;
+  raw: object | null;
+
+  constructor(
+    response: string,
+    sourceNodes?: NodeWithScore[],
+    chatResponse?: ChatResponse,
+  ) {
     this.response = response;
     this.sourceNodes = sourceNodes || [];
+    this.message = chatResponse?.message ?? {
+      content: response,
+      role: "assistant",
+    };
+    this.raw = chatResponse?.raw ?? null;
   }
 
   protected _getFormattedSources() {

packages/core/src/agent/anthropic.ts

@@ -1,19 +1,14 @@
+import { Response } from "../Response.js";
 import { Settings } from "../Settings.js";
 import {
   type ChatEngineParamsNonStreaming,
   type ChatEngineParamsStreaming,
 } from "../engines/chat/index.js";
 import { stringifyJSONToMessageContent } from "../internal/utils.js";
 import { Anthropic } from "../llm/anthropic.js";
-import type { ToolCallLLMMessageOptions } from "../llm/index.js";
 import { ObjectRetriever } from "../objects/index.js";
 import type { BaseToolWithCall } from "../types.js";
-import {
-  AgentRunner,
-  AgentWorker,
-  type AgentChatResponse,
-  type AgentParamsBase,
-} from "./base.js";
+import { AgentRunner, AgentWorker, type AgentParamsBase } from "./base.js";
 import type { TaskHandler } from "./types.js";
 import { callTool } from "./utils.js";
 
@@ -56,9 +51,7 @@ export class AnthropicAgent extends AgentRunner<Anthropic> {
 
   createStore = AgentRunner.defaultCreateStore;
 
-  async chat(
-    params: ChatEngineParamsNonStreaming,
-  ): Promise<AgentChatResponse<ToolCallLLMMessageOptions>>;
+  async chat(params: ChatEngineParamsNonStreaming): Promise<Response>;
   async chat(params: ChatEngineParamsStreaming): Promise<never>;
   override async chat(
     params: ChatEngineParamsNonStreaming | ChatEngineParamsStreaming,

packages/core/src/agent/base.ts

@@ -1,4 +1,5 @@
 import { ReadableStream, TransformStream, randomUUID } from "@llamaindex/env";
+import { Response } from "../Response.js";
 import { Settings } from "../Settings.js";
 import {
   type ChatEngine,
@@ -9,13 +10,8 @@ import { wrapEventCaller } from "../internal/context/EventCaller.js";
 import { consoleLogger, emptyLogger } from "../internal/logger.js";
 import { getCallbackManager } from "../internal/settings/CallbackManager.js";
 import { isAsyncIterable } from "../internal/utils.js";
-import type {
-  ChatMessage,
-  ChatResponse,
-  ChatResponseChunk,
-  LLM,
-  MessageContent,
-} from "../llm/index.js";
+import type { ChatMessage, LLM, MessageContent } from "../llm/index.js";
+import { extractText } from "../llm/utils.js";
 import type { BaseToolWithCall, ToolOutput } from "../types.js";
 import type {
   AgentTaskContext,
@@ -101,16 +97,6 @@ export function createTaskOutputStream<
   });
 }
 
-export type AgentStreamChatResponse<Options extends object> = {
-  response: ChatResponseChunk<Options>;
-  sources: ToolOutput[];
-};
-
-export type AgentChatResponse<Options extends object> = {
-  response: ChatResponse<Options>;
-  sources: ToolOutput[];
-};
-
 export type AgentRunnerParams<
   AI extends LLM,
   Store extends object = {},
@@ -210,11 +196,7 @@ export abstract class AgentRunner<
   >
     ? AdditionalMessageOptions
     : never,
-> implements
-    ChatEngine<
-      AgentChatResponse<AdditionalMessageOptions>,
-      ReadableStream<AgentStreamChatResponse<AdditionalMessageOptions>>
-    >
+> implements ChatEngine
 {
   readonly #llm: AI;
   readonly #tools:
@@ -320,47 +302,34 @@ export abstract class AgentRunner<
     });
   }
 
-  async chat(
-    params: ChatEngineParamsNonStreaming,
-  ): Promise<AgentChatResponse<AdditionalMessageOptions>>;
+  async chat(params: ChatEngineParamsNonStreaming): Promise<Response>;
   async chat(
     params: ChatEngineParamsStreaming,
-  ): Promise<ReadableStream<AgentStreamChatResponse<AdditionalMessageOptions>>>;
+  ): Promise<ReadableStream<Response>>;
   @wrapEventCaller
   async chat(
     params: ChatEngineParamsNonStreaming | ChatEngineParamsStreaming,
-  ): Promise<
-    | AgentChatResponse<AdditionalMessageOptions>
-    | ReadableStream<AgentStreamChatResponse<AdditionalMessageOptions>>
-  > {
+  ): Promise<Response | ReadableStream<Response>> {
     const task = this.createTask(params.message, !!params.stream);
     for await (const stepOutput of task) {
       // update chat history for each round
       this.#chatHistory = [...stepOutput.taskStep.context.store.messages];
       if (stepOutput.isLast) {
-        const { output, taskStep } = stepOutput;
+        const { output } = stepOutput;
         if (isAsyncIterable(output)) {
-          return output.pipeThrough<
-            AgentStreamChatResponse<AdditionalMessageOptions>
-          >(
+          return output.pipeThrough<Response>(
             new TransformStream({
               transform(chunk, controller) {
-                controller.enqueue({
-                  response: chunk,
-                  get sources() {
-                    return [...taskStep.context.store.toolOutputs];
-                  },
-                });
+                controller.enqueue(new Response(chunk.delta));
               },
             }),
           );
         } else {
-          return {
-            response: output,
-            get sources() {
-              return [...taskStep.context.store.toolOutputs];
-            },
-          } satisfies AgentChatResponse<AdditionalMessageOptions>;
+          return new Response(
+            extractText(output.message.content),
+            undefined,
+            output,
+          );
         }
       }
     }