support agent workflow

packages/server/src/handlers/chat.ts

-import { LlamaIndexAdapter } from "ai";
+import { type Message } from "ai";
 import { IncomingMessage, ServerResponse } from "http";
 import { type ChatMessage } from "llamaindex";
 import type { ServerWorkflow } from "../types";
@@ -7,7 +7,7 @@ import {
   pipeResponse,
   sendJSONResponse,
 } from "../utils/request";
-import { createStreamFromWorkflowContext } from "../utils/stream";
+import { runWorkflow } from "../utils/workflow";
 
 export const handleChat = async (
   workflow: ServerWorkflow,
@@ -16,7 +16,7 @@ export const handleChat = async (
 ) => {
   try {
     const body = await parseRequestBody(req);
-    const { messages } = body as { messages: ChatMessage[] };
+    const { messages } = body as { messages: Message[] };
 
     const lastMessage = messages[messages.length - 1];
     if (lastMessage?.role !== "user") {
@@ -25,15 +25,9 @@ export const handleChat = async (
       });
     }
 
-    const userMessage = lastMessage.content;
-    const chatHistory = messages.slice(0, -1);
-
-    const context = workflow.run({ userMessage, chatHistory });
-    const { stream, dataStream } =
-      await createStreamFromWorkflowContext(context);
-    const streamResponse = LlamaIndexAdapter.toDataStreamResponse(stream, {
-      data: dataStream,
-    });
+    const userInput = lastMessage.content;
+    const chatHistory = messages.slice(0, -1) as ChatMessage[];
+    const streamResponse = await runWorkflow(workflow, userInput, chatHistory);
     pipeResponse(res, streamResponse);
   } catch (error) {
     console.error("Chat error:", error);

packages/server/src/types.ts

 import {
+  AgentWorkflow,
   Workflow,
   type ChatMessage,
   type ChatResponseChunk,
-  type MessageContent,
 } from "llamaindex";
 
 export type AgentInput = {
-  userMessage: MessageContent;
-  chatHistory: ChatMessage[];
+  userInput: string; // the last message content from the user
+  chatHistory: ChatMessage[]; // the previous chat history (not including the last message)
 };
 
-export type ServerWorkflow = Workflow<null, AgentInput, ChatResponseChunk>;
+export type ServerWorkflow =
+  | Workflow<null, AgentInput, ChatResponseChunk>
+  | AgentWorkflow;

packages/server/src/utils/stream.ts → packages/server/src/utils/workflow.ts

-import { StreamData, type JSONValue } from "ai";
+import { LlamaIndexAdapter, StreamData, type JSONValue } from "ai";
 import {
+  AgentWorkflow,
   EngineResponse,
   StopEvent,
   WorkflowContext,
   WorkflowEvent,
+  type ChatMessage,
   type ChatResponseChunk,
 } from "llamaindex";
 import { ReadableStream } from "stream/web";
+import type { ServerWorkflow } from "../types";
 
-export async function createStreamFromWorkflowContext<Input, Output, Context>(
+export async function runWorkflow(
+  workflow: ServerWorkflow,
+  userInput: string,
+  chatHistory: ChatMessage[],
+) {
+  if (workflow instanceof AgentWorkflow) {
+    const context = workflow.run(userInput, { chatHistory });
+    const { stream, dataStream } = await createStreamFromWorkflowContext(
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      context as any,
+    );
+    return LlamaIndexAdapter.toDataStreamResponse(stream, { data: dataStream });
+  }
+
+  const context = workflow.run({ userInput, chatHistory });
+  const { stream, dataStream } = await createStreamFromWorkflowContext(context);
+  return LlamaIndexAdapter.toDataStreamResponse(stream, { data: dataStream });
+}
+
+async function createStreamFromWorkflowContext<Input, Output, Context>(
   context: WorkflowContext<Input, Output, Context>,
 ): Promise<{ stream: ReadableStream<EngineResponse>; dataStream: StreamData }> {
   const dataStream = new StreamData();