update for tools helper

---------
Co-authored-by: Marcus Schiesser <mail@marcusschiesser.de>

.changeset/itchy-peas-obey.md

@@ -0,0 +1,5 @@
+---
+"create-llama": patch
+---
+
+feat: bump LITS 0.8.2

.changeset/light-parrots-work.md

@@ -0,0 +1,5 @@
+---
+"create-llama": patch
+---
+
+Optimize generated workflow code for Python

.changeset/nice-garlics-repeat.md

@@ -0,0 +1,5 @@
+---
+"create-llama": patch
+---
+
+feat: use llamaindex chat-ui for nextjs frontend

e2e/shared/multiagent_template.spec.ts

@@ -18,7 +18,7 @@ const templateUI: TemplateUI = "shadcn";
 const templatePostInstallAction: TemplatePostInstallAction = "runApp";
 const appType: AppType = templateFramework === "nextjs" ? "" : "--frontend";
 const userMessage = "Write a blog post about physical standards for letters";
-const templateAgents = ["financial_report", "blog"];
+const templateAgents = ["financial_report", "blog", "form_filling"];
 
 for (const agents of templateAgents) {
   test.describe(`Test multiagent template ${agents} ${templateFramework} ${dataSource} ${templateUI} ${appType} ${templatePostInstallAction}`, async () => {
@@ -26,6 +26,10 @@ for (const agents of templateAgents) {
       process.platform !== "linux" || process.env.DATASOURCE === "--no-files",
       "The multiagent template currently only works with files. We also only run on Linux to speed up tests.",
     );
+    test.skip(
+      agents === "form_filling" && templateFramework !== "fastapi",
+      "Form filling is currently only supported with FastAPI.",
+    );
     let port: number;
     let externalPort: number;
     let cwd: string;
@@ -68,6 +72,10 @@ for (const agents of templateAgents) {
     test("Frontend should be able to submit a message and receive the start of a streamed response", async ({
       page,
     }) => {
+      test.skip(
+        agents === "financial_report" || agents === "form_filling",
+        "Skip chat tests for financial report and form filling.",
+      );
       await page.goto(`http://localhost:${port}`);
       await page.fill("form textarea", userMessage);
 

templates/components/agents/python/blog/README-template.md

@@ -8,9 +8,9 @@ This example is using three agents to generate a blog post:
 
 There are three different methods how the agents can interact to reach their goal:
 
-1. [Choreography](./app/examples/choreography.py) - the agents decide themselves to delegate a task to another agent
-1. [Orchestrator](./app/examples/orchestrator.py) - a central orchestrator decides which agent should execute a task
-1. [Explicit Workflow](./app/examples/workflow.py) - a pre-defined workflow specific for the task is used to execute the tasks
+1. [Choreography](./app/agents/choreography.py) - the agents decide themselves to delegate a task to another agent
+1. [Orchestrator](./app/agents/orchestrator.py) - a central orchestrator decides which agent should execute a task
+1. [Explicit Workflow](./app/agents/workflow.py) - a pre-defined workflow specific for the task is used to execute the tasks
 
 ## Getting Started
 

templates/components/agents/python/blog/app/workflows/__init__.py

@@ -0,0 +1,3 @@
+from .blog import create_workflow
+
+__all__ = ["create_workflow"]

templates/components/agents/python/blog/app/workflows/blog.py

@@ -4,17 +4,18 @@ from typing import List, Optional
 
 from app.agents.choreography import create_choreography
 from app.agents.orchestrator import create_orchestrator
-from app.agents.workflow import create_workflow
+from app.agents.workflow import create_workflow as create_blog_workflow
 from llama_index.core.chat_engine.types import ChatMessage
 from llama_index.core.workflow import Workflow
 
 logger = logging.getLogger("uvicorn")
 
 
-def get_chat_engine(
+def create_workflow(
     chat_history: Optional[List[ChatMessage]] = None, **kwargs
 ) -> Workflow:
-    # TODO: the EXAMPLE_TYPE could be passed as a chat config parameter?
+    # Chat filters are not supported yet
+    kwargs.pop("filters", None)
     agent_type = os.getenv("EXAMPLE_TYPE", "").lower()
     match agent_type:
         case "choreography":
@@ -22,7 +23,7 @@ def get_chat_engine(
         case "orchestrator":
             agent = create_orchestrator(chat_history, **kwargs)
         case _:
-            agent = create_workflow(chat_history, **kwargs)
+            agent = create_blog_workflow(chat_history, **kwargs)
 
     logger.info(f"Using agent pattern: {agent_type}")
 

templates/components/agents/python/blog/app/workflows/single.py

@@ -42,9 +42,9 @@ class AgentRunEvent(Event):
         return {
             "type": "agent",
             "data": {
-                "name": self.name,
+                "agent": self.name,
                 "type": self.event_type.value,
-                "msg": self.msg,
+                "text": self.msg,
                 "data": self.data,
             },
         }

templates/components/agents/python/financial_report/README-template.md

@@ -33,7 +33,7 @@ curl --location 'localhost:8000/api/chat' \
 --data '{ "messages": [{ "role": "user", "content": "Create a report comparing the finances of Apple and Tesla" }] }'
 ```
 
-You can start editing the API by modifying `app/api/routers/chat.py` or `app/financial_report/workflow.py`. The API auto-updates as you save the files.
+You can start editing the API by modifying `app/api/routers/chat.py` or `app/workflows/financial_report.py`. The API auto-updates as you save the files.
 
 Open [http://localhost:8000/docs](http://localhost:8000/docs) with your browser to see the Swagger UI of the API.
 

templates/components/agents/python/financial_report/app/agents/analyst.py

@@ -1,47 +0,0 @@
-from textwrap import dedent
-from typing import List, Tuple
-
-from app.engine.tools import ToolFactory
-from app.workflows.single import FunctionCallingAgent
-from llama_index.core.chat_engine.types import ChatMessage
-from llama_index.core.tools import FunctionTool
-
-
-def _get_analyst_params() -> Tuple[List[type[FunctionTool]], str, str]:
-    tools = []
-    prompt_instructions = dedent(
-        """
-        You are an expert in analyzing financial data.
-        You are given a task and a set of financial data to analyze. Your task is to analyze the financial data and return a report.
-        Your response should include a detailed analysis of the financial data, including any trends, patterns, or insights that you find.
-        Construct the analysis in a textual format like tables would be great!
-        Don't need to synthesize the data, just analyze and provide your findings.
-        Always use the provided information, don't make up any information yourself.
-        """
-    )
-    description = "Expert in analyzing financial data"
-    configured_tools = ToolFactory.from_env(map_result=True)
-    # Check if the interpreter tool is configured
-    if "interpret" in configured_tools.keys():
-        tools.append(configured_tools["interpret"])
-        prompt_instructions += dedent("""
-            You are able to visualize the financial data using code interpreter tool.
-            It's very useful to create and include visualizations to the report (make sure you include the right code and data for the visualization).
-            Never include any code into the report, just the visualization.
-        """)
-        description += (
-            ", able to visualize the financial data using code interpreter tool."
-        )
-    return tools, prompt_instructions, description
-
-
-def create_analyst(chat_history: List[ChatMessage]):
-    tools, prompt_instructions, description = _get_analyst_params()
-
-    return FunctionCallingAgent(
-        name="analyst",
-        tools=tools,
-        description=description,
-        system_prompt=dedent(prompt_instructions),
-        chat_history=chat_history,
-    )

templates/components/agents/python/financial_report/app/agents/reporter.py

@@ -1,44 +0,0 @@
-from textwrap import dedent
-from typing import List, Tuple
-
-from app.engine.tools import ToolFactory
-from app.workflows.single import FunctionCallingAgent
-from llama_index.core.chat_engine.types import ChatMessage
-from llama_index.core.tools import BaseTool
-
-
-def _get_reporter_params(
-    chat_history: List[ChatMessage],
-) -> Tuple[List[type[BaseTool]], str, str]:
-    tools: List[type[BaseTool]] = []
-    description = "Expert in representing a financial report"
-    prompt_instructions = dedent(
-        """
-        You are a report generation assistant tasked with producing a well-formatted report given parsed context.
-        Given a comprehensive analysis of the user request, your task is to synthesize the information and return a well-formatted report.
-
-        ## Instructions
-        You are responsible for representing the analysis in a well-formatted report. If tables or visualizations provided, add them to the right sections that are most relevant.
-        Use only the provided information to create the report. Do not make up any information yourself.
-        Finally, the report should be presented in markdown format.
-        """
-    )
-    configured_tools = ToolFactory.from_env(map_result=True)
-    if "generate_document" in configured_tools:  # type: ignore
-        tools.append(configured_tools["generate_document"])  # type: ignore
-        prompt_instructions += (
-            "\nYou are also able to generate a file document (PDF/HTML) of the report."
-        )
-        description += " and generate a file document (PDF/HTML) of the report."
-    return tools, description, prompt_instructions
-
-
-def create_reporter(chat_history: List[ChatMessage]):
-    tools, description, prompt_instructions = _get_reporter_params(chat_history)
-    return FunctionCallingAgent(
-        name="reporter",
-        tools=tools,
-        description=description,
-        system_prompt=prompt_instructions,
-        chat_history=chat_history,
-    )

templates/components/agents/python/financial_report/app/agents/researcher.py

@@ -1,105 +0,0 @@
-import os
-from textwrap import dedent
-from typing import List, Optional
-
-from app.engine.index import IndexConfig, get_index
-from app.workflows.single import FunctionCallingAgent
-from llama_index.core.chat_engine.types import ChatMessage
-from llama_index.core.tools import BaseTool, QueryEngineTool, ToolMetadata
-from llama_index.indices.managed.llama_cloud import LlamaCloudIndex
-
-
-def _create_query_engine_tools(params=None) -> Optional[list[type[BaseTool]]]:
-    """
-    Provide an agent worker that can be used to query the index.
-    """
-    # Add query tool if index exists
-    index_config = IndexConfig(**(params or {}))
-    index = get_index(index_config)
-    if index is None:
-        return None
-
-    top_k = int(os.getenv("TOP_K", 5))
-
-    # Construct query engine tools
-    tools = []
-    # If index is LlamaCloudIndex, we need to add chunk and doc retriever tools
-    if isinstance(index, LlamaCloudIndex):
-        # Document retriever
-        doc_retriever = index.as_query_engine(
-            retriever_mode="files_via_content",
-            similarity_top_k=top_k,
-        )
-        chunk_retriever = index.as_query_engine(
-            retriever_mode="chunks",
-            similarity_top_k=top_k,
-        )
-        tools.append(
-            QueryEngineTool(
-                query_engine=doc_retriever,
-                metadata=ToolMetadata(
-                    name="document_retriever",
-                    description=dedent(
-                        """
-                        Document retriever that retrieves entire documents from the corpus.
-                        ONLY use for research questions that may require searching over entire research reports.
-                        Will be slower and more expensive than chunk-level retrieval but may be necessary.
-                        """
-                    ),
-                ),
-            )
-        )
-        tools.append(
-            QueryEngineTool(
-                query_engine=chunk_retriever,
-                metadata=ToolMetadata(
-                    name="chunk_retriever",
-                    description=dedent(
-                        """
-                        Retrieves a small set of relevant document chunks from the corpus.
-                        Use for research questions that want to look up specific facts from the knowledge corpus,
-                        and need entire documents.
-                        """
-                    ),
-                ),
-            )
-        )
-    else:
-        query_engine = index.as_query_engine(
-            **({"similarity_top_k": top_k} if top_k != 0 else {})
-        )
-        tools.append(
-            QueryEngineTool(
-                query_engine=query_engine,
-                metadata=ToolMetadata(
-                    name="retrieve_information",
-                    description="Use this tool to retrieve information about the text corpus from the index.",
-                ),
-            )
-        )
-    return tools
-
-
-def create_researcher(chat_history: List[ChatMessage], **kwargs):
-    """
-    Researcher is an agent that take responsibility for using tools to complete a given task.
-    """
-    tools = _create_query_engine_tools(**kwargs)
-
-    if tools is None:
-        raise ValueError("No tools found for researcher agent")
-
-    return FunctionCallingAgent(
-        name="researcher",
-        tools=tools,
-        description="expert in retrieving any unknown content from the corpus",
-        system_prompt=dedent(
-            """
-            You are a researcher agent. You are responsible for retrieving information from the corpus.
-            ## Instructions
-            + Don't synthesize the information, just return the whole retrieved information.
-            + Don't need to retrieve the information that is already provided in the chat history and response with: "There is no new information, please reuse the information from the conversation."
-            """
-        ),
-        chat_history=chat_history,
-    )

templates/components/agents/python/financial_report/app/agents/workflow.py

@@ -1,177 +0,0 @@
-from textwrap import dedent
-from typing import AsyncGenerator, List, Optional
-
-from app.agents.analyst import create_analyst
-from app.agents.reporter import create_reporter
-from app.agents.researcher import create_researcher
-from app.workflows.single import AgentRunEvent, AgentRunResult, FunctionCallingAgent
-from llama_index.core.chat_engine.types import ChatMessage
-from llama_index.core.prompts import PromptTemplate
-from llama_index.core.settings import Settings
-from llama_index.core.workflow import (
-    Context,
-    Event,
-    StartEvent,
-    StopEvent,
-    Workflow,
-    step,
-)
-
-
-def create_workflow(chat_history: Optional[List[ChatMessage]] = None, **kwargs):
-    researcher = create_researcher(
-        chat_history=chat_history,
-        **kwargs,
-    )
-
-    analyst = create_analyst(chat_history=chat_history)
-
-    reporter = create_reporter(chat_history=chat_history)
-
-    workflow = FinancialReportWorkflow(timeout=360, chat_history=chat_history)
-
-    workflow.add_workflows(
-        researcher=researcher,
-        analyst=analyst,
-        reporter=reporter,
-    )
-    return workflow
-
-
-class ResearchEvent(Event):
-    input: str
-
-
-class AnalyzeEvent(Event):
-    input: str
-
-
-class ReportEvent(Event):
-    input: str
-
-
-class FinancialReportWorkflow(Workflow):
-    def __init__(
-        self, timeout: int = 360, chat_history: Optional[List[ChatMessage]] = None
-    ):
-        super().__init__(timeout=timeout)
-        self.chat_history = chat_history or []
-
-    @step()
-    async def start(self, ctx: Context, ev: StartEvent) -> ResearchEvent | ReportEvent:
-        # set streaming
-        ctx.data["streaming"] = getattr(ev, "streaming", False)
-        # start the workflow with researching about a topic
-        ctx.data["task"] = ev.input
-        ctx.data["user_input"] = ev.input
-
-        # Decision-making process
-        decision = await self._decide_workflow(ev.input, self.chat_history)
-
-        if decision != "publish":
-            return ResearchEvent(input=f"Research for this task: {ev.input}")
-        else:
-            chat_history_str = "\n".join(
-                [f"{msg.role}: {msg.content}" for msg in self.chat_history]
-            )
-            return ReportEvent(
-                input=f"Create a report based on the chat history\n{chat_history_str}\n\n and task: {ev.input}"
-            )
-
-    async def _decide_workflow(
-        self, input: str, chat_history: List[ChatMessage]
-    ) -> str:
-        # TODO: Refactor this by using prompt generation
-        prompt_template = PromptTemplate(
-            dedent(
-                """
-                You are an expert in decision-making, helping people create financial reports for the provided data.
-                If the user doesn't need to add or update anything, respond with 'publish'.
-                Otherwise, respond with 'research'.
-
-                Here is the chat history:
-                {chat_history}
-
-                The current user request is:
-                {input}
-
-                Given the chat history and the new user request, decide whether to create a report based on existing information.
-                Decision (respond with either 'not_publish' or 'publish'):
-            """
-            )
-        )
-
-        chat_history_str = "\n".join(
-            [f"{msg.role}: {msg.content}" for msg in chat_history]
-        )
-        prompt = prompt_template.format(chat_history=chat_history_str, input=input)
-
-        output = await Settings.llm.acomplete(prompt)
-        decision = output.text.strip().lower()
-
-        return "publish" if decision == "publish" else "research"
-
-    @step()
-    async def research(
-        self, ctx: Context, ev: ResearchEvent, researcher: FunctionCallingAgent
-    ) -> AnalyzeEvent:
-        result: AgentRunResult = await self.run_agent(ctx, researcher, ev.input)
-        content = result.response.message.content
-        return AnalyzeEvent(
-            input=dedent(
-                f"""
-                Given the following research content:
-                {content}
-                Provide a comprehensive analysis of the data for the user's request: {ctx.data["task"]}
-                """
-            )
-        )
-
-    @step()
-    async def analyze(
-        self, ctx: Context, ev: AnalyzeEvent, analyst: FunctionCallingAgent
-    ) -> ReportEvent | StopEvent:
-        result: AgentRunResult = await self.run_agent(ctx, analyst, ev.input)
-        content = result.response.message.content
-        return ReportEvent(
-            input=dedent(
-                f"""
-                Given the following analysis:
-                {content}
-                Create a report for the user's request: {ctx.data["task"]}
-                """
-            )
-        )
-
-    @step()
-    async def report(
-        self, ctx: Context, ev: ReportEvent, reporter: FunctionCallingAgent
-    ) -> StopEvent:
-        try:
-            result: AgentRunResult = await self.run_agent(
-                ctx, reporter, ev.input, streaming=ctx.data["streaming"]
-            )
-            return StopEvent(result=result)
-        except Exception as e:
-            ctx.write_event_to_stream(
-                AgentRunEvent(
-                    name=reporter.name,
-                    msg=f"Error creating a report: {e}",
-                )
-            )
-            return StopEvent(result=None)
-
-    async def run_agent(
-        self,
-        ctx: Context,
-        agent: FunctionCallingAgent,
-        input: str,
-        streaming: bool = False,
-    ) -> AgentRunResult | AsyncGenerator:
-        handler = agent.run(input=input, streaming=streaming)
-        # bubble all events while running the executor to the planner
-        async for event in handler.stream_events():
-            # Don't write the StopEvent from sub task to the stream
-            if type(event) is not StopEvent:
-                ctx.write_event_to_stream(event)
-        return await handler

templates/components/agents/python/financial_report/app/engine/engine.py

@@ -1,12 +0,0 @@
-from typing import List, Optional
-
-from app.agents.workflow import create_workflow
-from llama_index.core.chat_engine.types import ChatMessage
-from llama_index.core.workflow import Workflow
-
-
-def get_chat_engine(
-    chat_history: Optional[List[ChatMessage]] = None, **kwargs
-) -> Workflow:
-    agent_workflow = create_workflow(chat_history, **kwargs)
-    return agent_workflow

templates/components/agents/python/financial_report/app/workflows/__init__.py

@@ -0,0 +1,3 @@
+from .financial_report import create_workflow
+
+__all__ = ["create_workflow"]

templates/components/agents/python/financial_report/app/workflows/financial_report.py

@@ -0,0 +1,298 @@
+import os
+from typing import Any, Dict, List, Optional
+
+from app.engine.index import IndexConfig, get_index
+from app.engine.tools import ToolFactory
+from app.workflows.events import AgentRunEvent
+from app.workflows.tools import (
+    call_tools,
+    chat_with_tools,
+)
+from llama_index.core import Settings
+from llama_index.core.base.llms.types import ChatMessage, MessageRole
+from llama_index.core.indices.vector_store import VectorStoreIndex
+from llama_index.core.llms.function_calling import FunctionCallingLLM
+from llama_index.core.memory import ChatMemoryBuffer
+from llama_index.core.tools import FunctionTool, QueryEngineTool, ToolSelection
+from llama_index.core.workflow import (
+    Context,
+    Event,
+    StartEvent,
+    StopEvent,
+    Workflow,
+    step,
+)
+
+
+def create_workflow(
+    chat_history: Optional[List[ChatMessage]] = None,
+    params: Optional[Dict[str, Any]] = None,
+    filters: Optional[List[Any]] = None,
+) -> Workflow:
+    index_config = IndexConfig(**params)
+    index: VectorStoreIndex = get_index(config=index_config)
+    if index is None:
+        query_engine_tool = None
+    else:
+        top_k = int(os.getenv("TOP_K", 10))
+        query_engine = index.as_query_engine(similarity_top_k=top_k, filters=filters)
+        query_engine_tool = QueryEngineTool.from_defaults(query_engine=query_engine)
+
+    configured_tools: Dict[str, FunctionTool] = ToolFactory.from_env(map_result=True)  # type: ignore
+    code_interpreter_tool = configured_tools.get("interpret")
+    document_generator_tool = configured_tools.get("generate_document")
+
+    return FinancialReportWorkflow(
+        query_engine_tool=query_engine_tool,
+        code_interpreter_tool=code_interpreter_tool,
+        document_generator_tool=document_generator_tool,
+        chat_history=chat_history,
+    )
+
+
+class InputEvent(Event):
+    input: List[ChatMessage]
+    response: bool = False
+
+
+class ResearchEvent(Event):
+    input: list[ToolSelection]
+
+
+class AnalyzeEvent(Event):
+    input: list[ToolSelection] | ChatMessage
+
+
+class ReportEvent(Event):
+    input: list[ToolSelection]
+
+
+class FinancialReportWorkflow(Workflow):
+    """
+    A workflow to generate a financial report using indexed documents.
+
+    Requirements:
+    - Indexed documents containing financial data and a query engine tool to search them
+    - A code interpreter tool to analyze data and generate reports
+    - A document generator tool to create report files
+
+    Steps:
+    1. LLM Input: The LLM determines the next step based on function calling.
+       For example, if the model requests the query engine tool, it returns a ResearchEvent;
+       if it requests document generation, it returns a ReportEvent.
+    2. Research: Uses the query engine to find relevant chunks from indexed documents.
+       After gathering information, it requests analysis (step 3).
+    3. Analyze: Uses a custom prompt to analyze research results and can call the code
+       interpreter tool for visualization or calculation. Returns results to the LLM.
+    4. Report: Uses the document generator tool to create a report. Returns results to the LLM.
+    """
+
+    _default_system_prompt = """
+    You are a financial analyst who are given a set of tools to help you.
+    It's good to using appropriate tools for the user request and always use the information from the tools, don't make up anything yourself.
+    For the query engine tool, you should break down the user request into a list of queries and call the tool with the queries.
+    """
+
+    def __init__(
+        self,
+        query_engine_tool: QueryEngineTool,
+        code_interpreter_tool: FunctionTool,
+        document_generator_tool: FunctionTool,
+        llm: Optional[FunctionCallingLLM] = None,
+        timeout: int = 360,
+        chat_history: Optional[List[ChatMessage]] = None,
+        system_prompt: Optional[str] = None,
+    ):
+        super().__init__(timeout=timeout)
+        self.system_prompt = system_prompt or self._default_system_prompt
+        self.chat_history = chat_history or []
+        self.query_engine_tool = query_engine_tool
+        self.code_interpreter_tool = code_interpreter_tool
+        self.document_generator_tool = document_generator_tool
+        assert (
+            query_engine_tool is not None
+        ), "Query engine tool is not found. Try run generation script or upload a document file first."
+        assert code_interpreter_tool is not None, "Code interpreter tool is required"
+        assert (
+            document_generator_tool is not None
+        ), "Document generator tool is required"
+        self.tools = [
+            self.query_engine_tool,
+            self.code_interpreter_tool,
+            self.document_generator_tool,
+        ]
+        self.llm: FunctionCallingLLM = llm or Settings.llm
+        assert isinstance(self.llm, FunctionCallingLLM)
+        self.memory = ChatMemoryBuffer.from_defaults(
+            llm=self.llm, chat_history=self.chat_history
+        )
+
+    @step()
+    async def prepare_chat_history(self, ctx: Context, ev: StartEvent) -> InputEvent:
+        ctx.data["input"] = ev.input
+
+        if self.system_prompt:
+            system_msg = ChatMessage(
+                role=MessageRole.SYSTEM, content=self.system_prompt
+            )
+            self.memory.put(system_msg)
+
+        # Add user input to memory
+        self.memory.put(ChatMessage(role=MessageRole.USER, content=ev.input))
+
+        return InputEvent(input=self.memory.get())
+
+    @step()
+    async def handle_llm_input(  # type: ignore
+        self,
+        ctx: Context,
+        ev: InputEvent,
+    ) -> ResearchEvent | AnalyzeEvent | ReportEvent | StopEvent:
+        """
+        Handle an LLM input and decide the next step.
+        """
+        # Always use the latest chat history from the input
+        chat_history: list[ChatMessage] = ev.input
+
+        # Get tool calls
+        response = await chat_with_tools(
+            self.llm,
+            self.tools,  # type: ignore
+            chat_history,
+        )
+        if not response.has_tool_calls():
+            # If no tool call, return the response generator
+            return StopEvent(result=response.generator)
+        # calling different tools at the same time is not supported at the moment
+        # add an error message to tell the AI to process step by step
+        if response.is_calling_different_tools():
+            self.memory.put(
+                ChatMessage(
+                    role=MessageRole.ASSISTANT,
+                    content="Cannot call different tools at the same time. Try calling one tool at a time.",
+                )
+            )
+            return InputEvent(input=self.memory.get())
+        self.memory.put(response.tool_call_message)
+        match response.tool_name():
+            case self.code_interpreter_tool.metadata.name:
+                return AnalyzeEvent(input=response.tool_calls)
+            case self.document_generator_tool.metadata.name:
+                return ReportEvent(input=response.tool_calls)
+            case self.query_engine_tool.metadata.name:
+                return ResearchEvent(input=response.tool_calls)
+            case _:
+                raise ValueError(f"Unknown tool: {response.tool_name()}")
+
+    @step()
+    async def research(self, ctx: Context, ev: ResearchEvent) -> AnalyzeEvent:
+        """
+        Do a research to gather information for the user's request.
+        A researcher should have these tools: query engine, search engine, etc.
+        """
+        ctx.write_event_to_stream(
+            AgentRunEvent(
+                name="Researcher",
+                msg="Starting research",
+            )
+        )
+        tool_calls = ev.input
+
+        tool_messages = await call_tools(
+            ctx=ctx,
+            agent_name="Researcher",
+            tools=[self.query_engine_tool],
+            tool_calls=tool_calls,
+        )
+        self.memory.put_messages(tool_messages)
+        return AnalyzeEvent(
+            input=ChatMessage(
+                role=MessageRole.ASSISTANT,
+                content="I've finished the research. Please analyze the result.",
+            ),
+        )
+
+    @step()
+    async def analyze(self, ctx: Context, ev: AnalyzeEvent) -> InputEvent:
+        """
+        Analyze the research result.
+        """
+        ctx.write_event_to_stream(
+            AgentRunEvent(
+                name="Analyst",
+                msg="Starting analysis",
+            )
+        )
+        event_requested_by_workflow_llm = isinstance(ev.input, list)
+        # Requested by the workflow LLM Input step, it's a tool call
+        if event_requested_by_workflow_llm:
+            # Set the tool calls
+            tool_calls = ev.input
+        else:
+            # Otherwise, it's triggered by the research step
+            # Use a custom prompt and independent memory for the analyst agent
+            analysis_prompt = """
+            You are a financial analyst, you are given a research result and a set of tools to help you.
+            Always use the given information, don't make up anything yourself. If there is not enough information, you can asking for more information.
+            If you have enough numerical information, it's good to include some charts/visualizations to the report so you can use the code interpreter tool to generate a report.
+            """
+            # This is handled by analyst agent
+            # Clone the shared memory to avoid conflicting with the workflow.
+            chat_history = self.memory.get()
+            chat_history.append(
+                ChatMessage(role=MessageRole.SYSTEM, content=analysis_prompt)
+            )
+            chat_history.append(ev.input)  # type: ignore
+            # Check if the analyst agent needs to call tools
+            response = await chat_with_tools(
+                self.llm,
+                [self.code_interpreter_tool],
+                chat_history,
+            )
+            if not response.has_tool_calls():
+                # If no tool call, fallback analyst message to the workflow
+                analyst_msg = ChatMessage(
+                    role=MessageRole.ASSISTANT,
+                    content=await response.full_response(),
+                )
+                self.memory.put(analyst_msg)
+                return InputEvent(input=self.memory.get())
+            else:
+                # Set the tool calls and the tool call message to the memory
+                tool_calls = response.tool_calls
+                self.memory.put(response.tool_call_message)
+
+        # Call tools
+        tool_messages = await call_tools(
+            ctx=ctx,
+            agent_name="Analyst",
+            tools=[self.code_interpreter_tool],
+            tool_calls=tool_calls,  # type: ignore
+        )
+        self.memory.put_messages(tool_messages)
+
+        # Fallback to the input with the latest chat history
+        return InputEvent(input=self.memory.get())
+
+    @step()
+    async def report(self, ctx: Context, ev: ReportEvent) -> InputEvent:
+        """
+        Generate a report based on the analysis result.
+        """
+        ctx.write_event_to_stream(
+            AgentRunEvent(
+                name="Reporter",
+                msg="Starting report generation",
+            )
+        )
+        tool_calls = ev.input
+        tool_messages = await call_tools(
+            ctx=ctx,
+            agent_name="Reporter",
+            tools=[self.document_generator_tool],
+            tool_calls=tool_calls,
+        )
+        self.memory.put_messages(tool_messages)
+
+        # After the tool calls, fallback to the input with the latest chat history
+        return InputEvent(input=self.memory.get())

templates/components/agents/python/form_filling/README-template.md

@@ -39,7 +39,7 @@ curl --location 'localhost:8000/api/chat' \
 --data '{ "messages": [{ "role": "user", "content": "What can you do?" }] }'
 ```
 
-You can start editing the API by modifying `app/api/routers/chat.py` or `app/agents/form_filling.py`. The API auto-updates as you save the files.
+You can start editing the API by modifying `app/api/routers/chat.py` or `app/workflows/form_filling.py`. The API auto-updates as you save the files.
 
 Open [http://localhost:8000/docs](http://localhost:8000/docs) with your browser to see the Swagger UI of the API.
 

templates/components/agents/python/form_filling/app/agents/form_filling.py

@@ -1,397 +0,0 @@
-import os
-import uuid
-from enum import Enum
-from typing import AsyncGenerator, List, Optional
-
-from app.engine.index import get_index
-from app.engine.tools import ToolFactory
-from app.engine.tools.form_filling import CellValue, MissingCell
-from llama_index.core import Settings
-from llama_index.core.base.llms.types import ChatMessage, MessageRole
-from llama_index.core.indices.vector_store import VectorStoreIndex
-from llama_index.core.llms.function_calling import FunctionCallingLLM
-from llama_index.core.memory import ChatMemoryBuffer
-from llama_index.core.tools import FunctionTool, QueryEngineTool, ToolSelection
-from llama_index.core.tools.types import ToolOutput
-from llama_index.core.workflow import (
-    Context,
-    Event,
-    StartEvent,
-    StopEvent,
-    Workflow,
-    step,
-)
-from pydantic import Field
-
-
-def create_workflow(
-    chat_history: Optional[List[ChatMessage]] = None, **kwargs
-) -> Workflow:
-    index: VectorStoreIndex = get_index()
-    if index is None:
-        query_engine_tool = None
-    else:
-        top_k = int(os.getenv("TOP_K", 10))
-        query_engine = index.as_query_engine(similarity_top_k=top_k)
-        query_engine_tool = QueryEngineTool.from_defaults(query_engine=query_engine)
-
-    configured_tools = ToolFactory.from_env(map_result=True)
-    extractor_tool = configured_tools.get("extract_questions")
-    filling_tool = configured_tools.get("fill_form")
-
-    if extractor_tool is None or filling_tool is None:
-        raise ValueError("Extractor or filling tool is not found!")
-
-    workflow = FormFillingWorkflow(
-        query_engine_tool=query_engine_tool,
-        extractor_tool=extractor_tool,
-        filling_tool=filling_tool,
-        chat_history=chat_history,
-    )
-
-    return workflow
-
-
-class InputEvent(Event):
-    input: List[ChatMessage]
-    response: bool = False
-
-
-class ExtractMissingCellsEvent(Event):
-    tool_call: ToolSelection
-
-
-class FindAnswersEvent(Event):
-    missing_cells: list[MissingCell]
-
-
-class FillEvent(Event):
-    tool_call: ToolSelection
-
-
-class AgentRunEventType(Enum):
-    TEXT = "text"
-    PROGRESS = "progress"
-
-
-class AgentRunEvent(Event):
-    name: str
-    msg: str
-    event_type: AgentRunEventType = Field(default=AgentRunEventType.TEXT)
-    data: Optional[dict] = None
-
-    def to_response(self) -> dict:
-        return {
-            "type": "agent",
-            "data": {
-                "agent": self.name,
-                "type": self.event_type.value,
-                "text": self.msg,
-                "data": self.data,
-            },
-        }
-
-
-class FormFillingWorkflow(Workflow):
-    """
-    A predefined workflow for filling missing cells in a CSV file.
-    Required tools:
-    - query_engine: A query engine to query for the answers to the questions.
-    - extract_question: Extract missing cells in a CSV file and generate questions to fill them.
-    - answer_question: Query for the answers to the questions.
-
-    Flow:
-    1. Extract missing cells in a CSV file and generate questions to fill them.
-    2. Query for the answers to the questions.
-    3. Fill the missing cells with the answers.
-    """
-
-    _default_system_prompt = """
-    You are a helpful assistant who helps fill missing cells in a CSV file.
-    Only use provided data, never make up any information yourself. Fill N/A if the answer is not found.
-    """
-
-    def __init__(
-        self,
-        query_engine_tool: QueryEngineTool,
-        extractor_tool: FunctionTool,
-        filling_tool: FunctionTool,
-        llm: Optional[FunctionCallingLLM] = None,
-        timeout: int = 360,
-        chat_history: Optional[List[ChatMessage]] = None,
-        system_prompt: Optional[str] = None,
-    ):
-        super().__init__(timeout=timeout)
-        self.system_prompt = system_prompt or self._default_system_prompt
-        self.chat_history = chat_history or []
-        self.query_engine_tool = query_engine_tool
-        self.extractor_tool = extractor_tool
-        self.filling_tool = filling_tool
-        self.llm: FunctionCallingLLM = llm or Settings.llm
-        if not isinstance(self.llm, FunctionCallingLLM):
-            raise ValueError("FormFillingWorkflow only supports FunctionCallingLLM.")
-        self.memory = ChatMemoryBuffer.from_defaults(
-            llm=self.llm, chat_history=self.chat_history
-        )
-
-    @step()
-    async def start(self, ctx: Context, ev: StartEvent) -> InputEvent:
-        ctx.data["streaming"] = getattr(ev, "streaming", False)
-        ctx.data["input"] = ev.input
-
-        if self.system_prompt:
-            system_msg = ChatMessage(
-                role=MessageRole.SYSTEM, content=self.system_prompt
-            )
-            self.memory.put(system_msg)
-
-        user_input = ev.input
-        user_msg = ChatMessage(role=MessageRole.USER, content=user_input)
-        self.memory.put(user_msg)
-
-        chat_history = self.memory.get()
-        return InputEvent(input=chat_history)
-
-    @step(pass_context=True)
-    async def handle_llm_input(  # type: ignore
-        self,
-        ctx: Context,
-        ev: InputEvent,
-    ) -> ExtractMissingCellsEvent | FillEvent | StopEvent:
-        """
-        Handle an LLM input and decide the next step.
-        """
-        chat_history: list[ChatMessage] = ev.input
-
-        generator = self._tool_call_generator(chat_history)
-
-        # Check for immediate tool call
-        is_tool_call = await generator.__anext__()
-        if is_tool_call:
-            full_response = await generator.__anext__()
-            tool_calls = self.llm.get_tool_calls_from_response(full_response)  # type: ignore
-            for tool_call in tool_calls:
-                if tool_call.tool_name == self.extractor_tool.metadata.get_name():
-                    ctx.send_event(ExtractMissingCellsEvent(tool_call=tool_call))
-                elif tool_call.tool_name == self.filling_tool.metadata.get_name():
-                    ctx.send_event(FillEvent(tool_call=tool_call))
-        else:
-            # If no tool call, return the generator
-            return StopEvent(result=generator)
-
-    @step()
-    async def extract_missing_cells(
-        self, ctx: Context, ev: ExtractMissingCellsEvent
-    ) -> InputEvent | FindAnswersEvent:
-        """
-        Extract missing cells in a CSV file and generate questions to fill them.
-        """
-        ctx.write_event_to_stream(
-            AgentRunEvent(
-                name="Extractor",
-                msg="Extracting missing cells",
-            )
-        )
-        # Call the extract questions tool
-        response = self._call_tool(
-            ctx,
-            agent_name="Extractor",
-            tool=self.extractor_tool,
-            tool_selection=ev.tool_call,
-        )
-        if response.is_error:
-            return InputEvent(input=self.memory.get())
-
-        missing_cells = response.raw_output.get("missing_cells", [])
-        message = ChatMessage(
-            role=MessageRole.TOOL,
-            content=str(missing_cells),
-            additional_kwargs={
-                "tool_call_id": ev.tool_call.tool_id,
-                "name": ev.tool_call.tool_name,
-            },
-        )
-        self.memory.put(message)
-
-        if self.query_engine_tool is None:
-            # Fallback to input that query engine tool is not found so that cannot answer questions
-            self.memory.put(
-                ChatMessage(
-                    role=MessageRole.ASSISTANT,
-                    content="Extracted missing cells but query engine tool is not found so cannot answer questions. Ask user to upload file or connect to a knowledge base.",
-                )
-            )
-            return InputEvent(input=self.memory.get())
-
-        # Forward missing cells information to find answers step
-        return FindAnswersEvent(missing_cells=missing_cells)
-
-    @step()
-    async def find_answers(self, ctx: Context, ev: FindAnswersEvent) -> InputEvent:
-        """
-        Call answer questions tool to query for the answers to the questions.
-        """
-        ctx.write_event_to_stream(
-            AgentRunEvent(
-                name="Researcher",
-                msg="Finding answers for missing cells",
-            )
-        )
-        missing_cells = ev.missing_cells
-        # If missing cells information is not found, fallback to other tools
-        # It means that the extractor tool has not been called yet
-        # Fallback to input
-        if missing_cells is None:
-            ctx.write_event_to_stream(
-                AgentRunEvent(
-                    name="Researcher",
-                    msg="Error: Missing cells information not found. Fallback to other tools.",
-                )
-            )
-            message = ChatMessage(
-                role=MessageRole.TOOL,
-                content="Error: Missing cells information not found.",
-                additional_kwargs={
-                    "tool_call_id": ev.tool_call.tool_id,
-                    "name": ev.tool_call.tool_name,
-                },
-            )
-            self.memory.put(message)
-            return InputEvent(input=self.memory.get())
-
-        cell_values: list[CellValue] = []
-        # Iterate over missing cells and query for the answers
-        # and stream the progress
-        progress_id = str(uuid.uuid4())
-        total_steps = len(missing_cells)
-        for i, cell in enumerate(missing_cells):
-            if cell.question_to_answer is None:
-                continue
-            ctx.write_event_to_stream(
-                AgentRunEvent(
-                    name="Researcher",
-                    msg=f"Querying for: {cell.question_to_answer}",
-                    event_type=AgentRunEventType.PROGRESS,
-                    data={
-                        "id": progress_id,
-                        "total": total_steps,
-                        "current": i,
-                    },
-                )
-            )
-            # Call query engine tool directly
-            answer = await self.query_engine_tool.acall(query=cell.question_to_answer)
-            cell_values.append(
-                CellValue(
-                    row_index=cell.row_index,
-                    column_index=cell.column_index,
-                    value=str(answer),
-                )
-            )
-        self.memory.put(
-            ChatMessage(
-                role=MessageRole.ASSISTANT,
-                content=str(cell_values),
-            )
-        )
-        return InputEvent(input=self.memory.get())
-
-    @step()
-    async def fill_cells(self, ctx: Context, ev: FillEvent) -> InputEvent:
-        """
-        Call fill cells tool to fill the missing cells with the answers.
-        """
-        ctx.write_event_to_stream(
-            AgentRunEvent(
-                name="Processor",
-                msg="Filling missing cells",
-            )
-        )
-        # Call the fill cells tool
-        result = self._call_tool(
-            ctx,
-            agent_name="Processor",
-            tool=self.filling_tool,
-            tool_selection=ev.tool_call,
-        )
-        if result.is_error:
-            return InputEvent(input=self.memory.get())
-
-        message = ChatMessage(
-            role=MessageRole.TOOL,
-            content=str(result.raw_output),
-            additional_kwargs={
-                "tool_call_id": ev.tool_call.tool_id,
-                "name": ev.tool_call.tool_name,
-            },
-        )
-        self.memory.put(message)
-        return InputEvent(input=self.memory.get(), response=True)
-
-    async def _tool_call_generator(
-        self, chat_history: list[ChatMessage]
-    ) -> AsyncGenerator[ChatMessage | bool, None]:
-        response_stream = await self.llm.astream_chat_with_tools(
-            [self.extractor_tool, self.filling_tool],
-            chat_history=chat_history,
-        )
-
-        full_response = None
-        yielded_indicator = False
-        async for chunk in response_stream:
-            if "tool_calls" not in chunk.message.additional_kwargs:
-                # Yield a boolean to indicate whether the response is a tool call
-                if not yielded_indicator:
-                    yield False
-                    yielded_indicator = True
-
-                # if not a tool call, yield the chunks!
-                yield chunk
-            elif not yielded_indicator:
-                # Yield the indicator for a tool call
-                yield True
-                yielded_indicator = True
-
-            full_response = chunk
-
-        # Write the full response to memory and yield it
-        if full_response:
-            self.memory.put(full_response.message)
-            yield full_response
-
-    def _call_tool(
-        self,
-        ctx: Context,
-        agent_name: str,
-        tool: FunctionTool,
-        tool_selection: ToolSelection,
-    ) -> ToolOutput:
-        """
-        Safely call a tool and handle errors.
-        """
-        try:
-            response: ToolOutput = tool.call(**tool_selection.tool_kwargs)
-            return response
-        except Exception as e:
-            ctx.write_event_to_stream(
-                AgentRunEvent(
-                    name=agent_name,
-                    msg=f"Error: {str(e)}",
-                )
-            )
-            message = ChatMessage(
-                role=MessageRole.TOOL,
-                content=f"Error: {str(e)}",
-                additional_kwargs={
-                    "tool_call_id": tool_selection.tool_id,
-                    "name": tool.metadata.get_name(),
-                },
-            )
-            self.memory.put(message)
-            return ToolOutput(
-                content=f"Error: {str(e)}",
-                tool_name=tool.metadata.get_name(),
-                raw_input=tool_selection.tool_kwargs,
-                raw_output=None,
-                is_error=True,
-            )

templates/components/agents/python/form_filling/app/engine/engine.py

@@ -1,11 +0,0 @@
-from typing import List, Optional
-
-from app.agents.form_filling import create_workflow
-from llama_index.core.chat_engine.types import ChatMessage
-from llama_index.core.workflow import Workflow
-
-
-def get_chat_engine(
-    chat_history: Optional[List[ChatMessage]] = None, **kwargs
-) -> Workflow:
-    return create_workflow(chat_history=chat_history, **kwargs)

templates/components/agents/python/form_filling/app/workflows/__init__.py

@@ -0,0 +1,3 @@
+from .form_filling import create_workflow
+
+__all__ = ["create_workflow"]

templates/components/agents/python/form_filling/app/workflows/form_filling.py

@@ -0,0 +1,241 @@
+import os
+from typing import Any, Dict, List, Optional
+
+from app.engine.index import IndexConfig, get_index
+from app.engine.tools import ToolFactory
+from app.workflows.events import AgentRunEvent
+from app.workflows.tools import (
+    call_tools,
+    chat_with_tools,
+)
+from llama_index.core import Settings
+from llama_index.core.base.llms.types import ChatMessage, MessageRole
+from llama_index.core.indices.vector_store import VectorStoreIndex
+from llama_index.core.llms.function_calling import FunctionCallingLLM
+from llama_index.core.memory import ChatMemoryBuffer
+from llama_index.core.tools import FunctionTool, QueryEngineTool, ToolSelection
+from llama_index.core.workflow import (
+    Context,
+    Event,
+    StartEvent,
+    StopEvent,
+    Workflow,
+    step,
+)
+
+
+def create_workflow(
+    chat_history: Optional[List[ChatMessage]] = None,
+    params: Optional[Dict[str, Any]] = None,
+    filters: Optional[List[Any]] = None,
+) -> Workflow:
+    if params is None:
+        params = {}
+    if filters is None:
+        filters = []
+    index_config = IndexConfig(**params)
+    index: VectorStoreIndex = get_index(config=index_config)
+    if index is None:
+        query_engine_tool = None
+    else:
+        top_k = int(os.getenv("TOP_K", 10))
+        query_engine = index.as_query_engine(similarity_top_k=top_k, filters=filters)
+        query_engine_tool = QueryEngineTool.from_defaults(query_engine=query_engine)
+
+    configured_tools = ToolFactory.from_env(map_result=True)
+    extractor_tool = configured_tools.get("extract_questions")  # type: ignore
+    filling_tool = configured_tools.get("fill_form")  # type: ignore
+
+    workflow = FormFillingWorkflow(
+        query_engine_tool=query_engine_tool,
+        extractor_tool=extractor_tool,  # type: ignore
+        filling_tool=filling_tool,  # type: ignore
+        chat_history=chat_history,
+    )
+
+    return workflow
+
+
+class InputEvent(Event):
+    input: List[ChatMessage]
+    response: bool = False
+
+
+class ExtractMissingCellsEvent(Event):
+    tool_calls: list[ToolSelection]
+
+
+class FindAnswersEvent(Event):
+    tool_calls: list[ToolSelection]
+
+
+class FillEvent(Event):
+    tool_calls: list[ToolSelection]
+
+
+class FormFillingWorkflow(Workflow):
+    """
+    A predefined workflow for filling missing cells in a CSV file.
+    Required tools:
+    - query_engine: A query engine to query for the answers to the questions.
+    - extract_question: Extract missing cells in a CSV file and generate questions to fill them.
+    - answer_question: Query for the answers to the questions.
+
+    Flow:
+    1. Extract missing cells in a CSV file and generate questions to fill them.
+    2. Query for the answers to the questions.
+    3. Fill the missing cells with the answers.
+    """
+
+    _default_system_prompt = """
+    You are a helpful assistant who helps fill missing cells in a CSV file.
+    Only extract missing cells from CSV files.
+    Only use provided data - never make up any information yourself. Fill N/A if an answer is not found.
+    If there is no query engine tool or the gathered information has many N/A values indicating the questions don't match the data, respond with a warning and ask the user to upload a different file or connect to a knowledge base.
+    """
+
+    def __init__(
+        self,
+        query_engine_tool: Optional[QueryEngineTool],
+        extractor_tool: FunctionTool,
+        filling_tool: FunctionTool,
+        llm: Optional[FunctionCallingLLM] = None,
+        timeout: int = 360,
+        chat_history: Optional[List[ChatMessage]] = None,
+        system_prompt: Optional[str] = None,
+    ):
+        super().__init__(timeout=timeout)
+        self.system_prompt = system_prompt or self._default_system_prompt
+        self.chat_history = chat_history or []
+        self.query_engine_tool = query_engine_tool
+        self.extractor_tool = extractor_tool
+        self.filling_tool = filling_tool
+        if self.extractor_tool is None or self.filling_tool is None:
+            raise ValueError("Extractor and filling tools are required.")
+        self.tools = [self.extractor_tool, self.filling_tool]
+        if self.query_engine_tool is not None:
+            self.tools.append(self.query_engine_tool)  # type: ignore
+        self.llm: FunctionCallingLLM = llm or Settings.llm
+        if not isinstance(self.llm, FunctionCallingLLM):
+            raise ValueError("FormFillingWorkflow only supports FunctionCallingLLM.")
+        self.memory = ChatMemoryBuffer.from_defaults(
+            llm=self.llm, chat_history=self.chat_history
+        )
+
+    @step()
+    async def start(self, ctx: Context, ev: StartEvent) -> InputEvent:
+        ctx.data["input"] = ev.input
+
+        if self.system_prompt:
+            system_msg = ChatMessage(
+                role=MessageRole.SYSTEM, content=self.system_prompt
+            )
+            self.memory.put(system_msg)
+
+        user_input = ev.input
+        user_msg = ChatMessage(role=MessageRole.USER, content=user_input)
+        self.memory.put(user_msg)
+
+        chat_history = self.memory.get()
+        return InputEvent(input=chat_history)
+
+    @step()
+    async def handle_llm_input(  # type: ignore
+        self,
+        ctx: Context,
+        ev: InputEvent,
+    ) -> ExtractMissingCellsEvent | FillEvent | StopEvent:
+        """
+        Handle an LLM input and decide the next step.
+        """
+        chat_history: list[ChatMessage] = ev.input
+        response = await chat_with_tools(
+            self.llm,
+            self.tools,
+            chat_history,
+        )
+        if not response.has_tool_calls():
+            return StopEvent(result=response.generator)
+        # calling different tools at the same time is not supported at the moment
+        # add an error message to tell the AI to process step by step
+        if response.is_calling_different_tools():
+            self.memory.put(
+                ChatMessage(
+                    role=MessageRole.ASSISTANT,
+                    content="Cannot call different tools at the same time. Try calling one tool at a time.",
+                )
+            )
+            return InputEvent(input=self.memory.get())
+        self.memory.put(response.tool_call_message)
+        match response.tool_name():
+            case self.extractor_tool.metadata.name:
+                return ExtractMissingCellsEvent(tool_calls=response.tool_calls)
+            case self.query_engine_tool.metadata.name:
+                return FindAnswersEvent(tool_calls=response.tool_calls)
+            case self.filling_tool.metadata.name:
+                return FillEvent(tool_calls=response.tool_calls)
+            case _:
+                raise ValueError(f"Unknown tool: {response.tool_name()}")
+
+    @step()
+    async def extract_missing_cells(
+        self, ctx: Context, ev: ExtractMissingCellsEvent
+    ) -> InputEvent | FindAnswersEvent:
+        """
+        Extract missing cells in a CSV file and generate questions to fill them.
+        """
+        ctx.write_event_to_stream(
+            AgentRunEvent(
+                name="Extractor",
+                msg="Extracting missing cells",
+            )
+        )
+        # Call the extract questions tool
+        tool_messages = await call_tools(
+            agent_name="Extractor",
+            tools=[self.extractor_tool],
+            ctx=ctx,
+            tool_calls=ev.tool_calls,
+        )
+        self.memory.put_messages(tool_messages)
+        return InputEvent(input=self.memory.get())
+
+    @step()
+    async def find_answers(self, ctx: Context, ev: FindAnswersEvent) -> InputEvent:
+        """
+        Call answer questions tool to query for the answers to the questions.
+        """
+        ctx.write_event_to_stream(
+            AgentRunEvent(
+                name="Researcher",
+                msg="Finding answers for missing cells",
+            )
+        )
+        tool_messages = await call_tools(
+            ctx=ctx,
+            agent_name="Researcher",
+            tools=[self.query_engine_tool],
+            tool_calls=ev.tool_calls,
+        )
+        self.memory.put_messages(tool_messages)
+        return InputEvent(input=self.memory.get())
+
+    @step()
+    async def fill_cells(self, ctx: Context, ev: FillEvent) -> InputEvent:
+        """
+        Call fill cells tool to fill the missing cells with the answers.
+        """
+        ctx.write_event_to_stream(
+            AgentRunEvent(
+                name="Processor",
+                msg="Filling missing cells",
+            )
+        )
+        tool_messages = await call_tools(
+            agent_name="Processor",
+            tools=[self.filling_tool],
+            ctx=ctx,
+            tool_calls=ev.tool_calls,
+        )
+        self.memory.put_messages(tool_messages)
+        return InputEvent(input=self.memory.get())

templates/components/loaders/typescript/file/loader.ts

@@ -1,7 +1,7 @@
 import {
   FILE_EXT_TO_READER,
   SimpleDirectoryReader,
-} from "llamaindex/readers/SimpleDirectoryReader";
+} from "llamaindex/readers/index";
 
 export const DATA_DIR = "./data";
 

templates/components/loaders/typescript/llama_parse/loader.ts

@@ -2,7 +2,7 @@ import { LlamaParseReader } from "llamaindex";
 import {
   FILE_EXT_TO_READER,
   SimpleDirectoryReader,
-} from "llamaindex/readers/SimpleDirectoryReader";
+} from "llamaindex/readers/index";
 
 export const DATA_DIR = "./data";
 

templates/components/multiagent/python/app/api/routers/chat.py

@@ -4,7 +4,8 @@ from app.api.routers.models import (
     ChatData,
 )
 from app.api.routers.vercel_response import VercelStreamResponse
-from app.engine.engine import get_chat_engine
+from app.engine.query_filter import generate_filters
+from app.workflows import create_workflow
 from fastapi import APIRouter, BackgroundTasks, HTTPException, Request, status
 
 chat_router = r = APIRouter()
@@ -22,19 +23,20 @@ async def chat(
         last_message_content = data.get_last_message_content()
         messages = data.get_history_messages(include_agent_messages=True)
 
-        # The chat API supports passing private document filters and chat params
-        # but agent workflow does not support them yet
-        # ignore chat params and use all documents for now
-        # TODO: generate filters based on doc_ids
+        doc_ids = data.get_chat_document_ids()
+        filters = generate_filters(doc_ids)
         params = data.data or {}
-        engine = get_chat_engine(chat_history=messages, params=params)
 
-        event_handler = engine.run(input=last_message_content, streaming=True)
+        workflow = create_workflow(
+            chat_history=messages, params=params, filters=filters
+        )
+
+        event_handler = workflow.run(input=last_message_content, streaming=True)
         return VercelStreamResponse(
             request=request,
             chat_data=data,
             event_handler=event_handler,
-            events=engine.stream_events(),
+            events=workflow.stream_events(),
         )
     except Exception as e:
         logger.exception("Error in chat engine", exc_info=True)

templates/components/multiagent/python/app/workflows/events.py

@@ -0,0 +1,27 @@
+from enum import Enum
+from typing import Optional
+
+from llama_index.core.workflow import Event
+
+
+class AgentRunEventType(Enum):
+    TEXT = "text"
+    PROGRESS = "progress"
+
+
+class AgentRunEvent(Event):
+    name: str
+    msg: str
+    event_type: AgentRunEventType = AgentRunEventType.TEXT
+    data: Optional[dict] = None
+
+    def to_response(self) -> dict:
+        return {
+            "type": "agent",
+            "data": {
+                "agent": self.name,
+                "type": self.event_type.value,
+                "text": self.msg,
+                "data": self.data,
+            },
+        }

templates/components/multiagent/python/app/workflows/function_calling_agent.py

@@ -0,0 +1,121 @@
+from typing import Any, List, Optional
+
+from app.workflows.events import AgentRunEvent
+from app.workflows.tools import ToolCallResponse, call_tools, chat_with_tools
+from llama_index.core.base.llms.types import ChatMessage
+from llama_index.core.llms.function_calling import FunctionCallingLLM
+from llama_index.core.memory import ChatMemoryBuffer
+from llama_index.core.settings import Settings
+from llama_index.core.tools.types import BaseTool
+from llama_index.core.workflow import (
+    Context,
+    Event,
+    StartEvent,
+    StopEvent,
+    Workflow,
+    step,
+)
+
+
+class InputEvent(Event):
+    input: list[ChatMessage]
+
+
+class ToolCallEvent(Event):
+    input: ToolCallResponse
+
+
+class FunctionCallingAgent(Workflow):
+    """
+    A simple workflow to request LLM with tools independently.
+    You can share the previous chat history to provide the context for the LLM.
+    """
+
+    def __init__(
+        self,
+        *args: Any,
+        llm: FunctionCallingLLM | None = None,
+        chat_history: Optional[List[ChatMessage]] = None,
+        tools: List[BaseTool] | None = None,
+        system_prompt: str | None = None,
+        verbose: bool = False,
+        timeout: float = 360.0,
+        name: str,
+        write_events: bool = True,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(*args, verbose=verbose, timeout=timeout, **kwargs)  # type: ignore
+        self.tools = tools or []
+        self.name = name
+        self.write_events = write_events
+
+        if llm is None:
+            llm = Settings.llm
+        self.llm = llm
+        if not self.llm.metadata.is_function_calling_model:
+            raise ValueError("The provided LLM must support function calling.")
+
+        self.system_prompt = system_prompt
+
+        self.memory = ChatMemoryBuffer.from_defaults(
+            llm=self.llm, chat_history=chat_history
+        )
+        self.sources = []  # type: ignore
+
+    @step()
+    async def prepare_chat_history(self, ctx: Context, ev: StartEvent) -> InputEvent:
+        # clear sources
+        self.sources = []
+
+        # set streaming
+        ctx.data["streaming"] = getattr(ev, "streaming", False)
+
+        # set system prompt
+        if self.system_prompt is not None:
+            system_msg = ChatMessage(role="system", content=self.system_prompt)
+            self.memory.put(system_msg)
+
+        # get user input
+        user_input = ev.input
+        user_msg = ChatMessage(role="user", content=user_input)
+        self.memory.put(user_msg)
+
+        if self.write_events:
+            ctx.write_event_to_stream(
+                AgentRunEvent(name=self.name, msg=f"Start to work on: {user_input}")
+            )
+
+        return InputEvent(input=self.memory.get())
+
+    @step()
+    async def handle_llm_input(
+        self,
+        ctx: Context,
+        ev: InputEvent,
+    ) -> ToolCallEvent | StopEvent:
+        chat_history = ev.input
+
+        response = await chat_with_tools(
+            self.llm,
+            self.tools,
+            chat_history,
+        )
+        is_tool_call = isinstance(response, ToolCallResponse)
+        if not is_tool_call:
+            if ctx.data["streaming"]:
+                return StopEvent(result=response)
+            else:
+                full_response = ""
+                async for chunk in response.generator:
+                    full_response += chunk.message.content
+                return StopEvent(result=full_response)
+        return ToolCallEvent(input=response)
+
+    @step()
+    async def handle_tool_calls(self, ctx: Context, ev: ToolCallEvent) -> InputEvent:
+        tool_calls = ev.input.tool_calls
+        tool_call_message = ev.input.tool_call_message
+        self.memory.put(tool_call_message)
+        tool_messages = await call_tools(self.name, self.tools, ctx, tool_calls)
+        self.memory.put_messages(tool_messages)
+        return InputEvent(input=self.memory.get())

templates/components/multiagent/python/app/workflows/tools.py

@@ -0,0 +1,237 @@
+import logging
+import uuid
+from abc import ABC, abstractmethod
+from typing import Any, AsyncGenerator, Callable, Optional
+
+from app.workflows.events import AgentRunEvent, AgentRunEventType
+from llama_index.core.base.llms.types import ChatMessage, ChatResponse, MessageRole
+from llama_index.core.llms.function_calling import FunctionCallingLLM
+from llama_index.core.tools import (
+    BaseTool,
+    FunctionTool,
+    ToolOutput,
+    ToolSelection,
+)
+from llama_index.core.workflow import Context
+from pydantic import BaseModel, ConfigDict
+
+logger = logging.getLogger("uvicorn")
+
+
+class ContextAwareTool(FunctionTool, ABC):
+    @abstractmethod
+    async def acall(self, ctx: Context, input: Any) -> ToolOutput:  # type: ignore
+        pass
+
+
+class ResponseGenerator(BaseModel):
+    """
+    A response generator from chat_with_tools.
+    """
+
+    generator: AsyncGenerator[ChatResponse | None, None]
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+
+class ChatWithToolsResponse(BaseModel):
+    """
+    A tool call response from chat_with_tools.
+    """
+
+    tool_calls: Optional[list[ToolSelection]]
+    tool_call_message: Optional[ChatMessage]
+    generator: Optional[AsyncGenerator[ChatResponse | None, None]]
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    def is_calling_different_tools(self) -> bool:
+        tool_names = {tool_call.tool_name for tool_call in self.tool_calls}
+        return len(tool_names) > 1
+
+    def has_tool_calls(self) -> bool:
+        return self.tool_calls is not None and len(self.tool_calls) > 0
+
+    def tool_name(self) -> str:
+        assert self.has_tool_calls()
+        assert not self.is_calling_different_tools()
+        return self.tool_calls[0].tool_name
+
+    async def full_response(self) -> str:
+        assert self.generator is not None
+        full_response = ""
+        async for chunk in self.generator:
+            full_response += chunk.message.content
+        return full_response
+
+
+async def chat_with_tools(  # type: ignore
+    llm: FunctionCallingLLM,
+    tools: list[BaseTool],
+    chat_history: list[ChatMessage],
+) -> ChatWithToolsResponse:
+    """
+    Request LLM to call tools or not.
+    This function doesn't change the memory.
+    """
+    generator = _tool_call_generator(llm, tools, chat_history)
+    is_tool_call = await generator.__anext__()
+    if is_tool_call:
+        # Last chunk is the full response
+        # Wait for the last chunk
+        full_response = None
+        async for chunk in generator:
+            full_response = chunk
+        assert isinstance(full_response, ChatResponse)
+        return ChatWithToolsResponse(
+            tool_calls=llm.get_tool_calls_from_response(full_response),
+            tool_call_message=full_response.message,
+            generator=None,
+        )
+    else:
+        return ChatWithToolsResponse(
+            tool_calls=None,
+            tool_call_message=None,
+            generator=generator,
+        )
+
+
+async def call_tools(
+    ctx: Context,
+    agent_name: str,
+    tools: list[BaseTool],
+    tool_calls: list[ToolSelection],
+    emit_agent_events: bool = True,
+) -> list[ChatMessage]:
+    if len(tool_calls) == 0:
+        return []
+
+    tools_by_name = {tool.metadata.get_name(): tool for tool in tools}
+    if len(tool_calls) == 1:
+        return [
+            await call_tool(
+                ctx,
+                tools_by_name[tool_calls[0].tool_name],
+                tool_calls[0],
+                lambda msg: ctx.write_event_to_stream(
+                    AgentRunEvent(
+                        name=agent_name,
+                        msg=msg,
+                    )
+                ),
+            )
+        ]
+    # Multiple tool calls, show progress
+    tool_msgs: list[ChatMessage] = []
+
+    progress_id = str(uuid.uuid4())
+    total_steps = len(tool_calls)
+    if emit_agent_events:
+        ctx.write_event_to_stream(
+            AgentRunEvent(
+                name=agent_name,
+                msg=f"Making {total_steps} tool calls",
+            )
+        )
+    for i, tool_call in enumerate(tool_calls):
+        tool = tools_by_name.get(tool_call.tool_name)
+        if not tool:
+            tool_msgs.append(
+                ChatMessage(
+                    role=MessageRole.ASSISTANT,
+                    content=f"Tool {tool_call.tool_name} does not exist",
+                )
+            )
+            continue
+        tool_msg = await call_tool(
+            ctx,
+            tool,
+            tool_call,
+            event_emitter=lambda msg: ctx.write_event_to_stream(
+                AgentRunEvent(
+                    name=agent_name,
+                    msg=msg,
+                    event_type=AgentRunEventType.PROGRESS,
+                    data={
+                        "id": progress_id,
+                        "total": total_steps,
+                        "current": i,
+                    },
+                )
+            ),
+        )
+        tool_msgs.append(tool_msg)
+    return tool_msgs
+
+
+async def call_tool(
+    ctx: Context,
+    tool: BaseTool,
+    tool_call: ToolSelection,
+    event_emitter: Optional[Callable[[str], None]],
+) -> ChatMessage:
+    if event_emitter:
+        event_emitter(
+            f"Calling tool {tool_call.tool_name}, {str(tool_call.tool_kwargs)}"
+        )
+    try:
+        if isinstance(tool, ContextAwareTool):
+            if ctx is None:
+                raise ValueError("Context is required for context aware tool")
+            # inject context for calling an context aware tool
+            response = await tool.acall(ctx=ctx, **tool_call.tool_kwargs)
+        else:
+            response = await tool.acall(**tool_call.tool_kwargs)  # type: ignore
+        return ChatMessage(
+            role=MessageRole.TOOL,
+            content=str(response.raw_output),
+            additional_kwargs={
+                "tool_call_id": tool_call.tool_id,
+                "name": tool.metadata.get_name(),
+            },
+        )
+    except Exception as e:
+        logger.error(f"Got error in tool {tool_call.tool_name}: {str(e)}")
+        if event_emitter:
+            event_emitter(f"Got error in tool {tool_call.tool_name}: {str(e)}")
+        return ChatMessage(
+            role=MessageRole.TOOL,
+            content=f"Error: {str(e)}",
+            additional_kwargs={
+                "tool_call_id": tool_call.tool_id,
+                "name": tool.metadata.get_name(),
+            },
+        )
+
+
+async def _tool_call_generator(
+    llm: FunctionCallingLLM,
+    tools: list[BaseTool],
+    chat_history: list[ChatMessage],
+) -> AsyncGenerator[ChatResponse | bool, None]:
+    response_stream = await llm.astream_chat_with_tools(
+        tools,
+        chat_history=chat_history,
+        allow_parallel_tool_calls=False,
+    )
+
+    full_response = None
+    yielded_indicator = False
+    async for chunk in response_stream:
+        if "tool_calls" not in chunk.message.additional_kwargs:
+            # Yield a boolean to indicate whether the response is a tool call
+            if not yielded_indicator:
+                yield False
+                yielded_indicator = True
+
+            # if not a tool call, yield the chunks!
+            yield chunk  # type: ignore
+        elif not yielded_indicator:
+            # Yield the indicator for a tool call
+            yield True
+            yielded_indicator = True
+
+        full_response = chunk
+
+    if full_response:
+        yield full_response  # type: ignore

templates/types/streaming/express/package.json

@@ -20,7 +20,7 @@
     "dotenv": "^16.3.1",
     "duck-duck-scrape": "^2.2.5",
     "express": "^4.18.2",
-    "llamaindex": "0.7.10",
+    "llamaindex": "0.8.2",
     "pdf2json": "3.0.5",
     "ajv": "^8.12.0",
     "@e2b/code-interpreter": "0.0.9-beta.3",

templates/types/streaming/nextjs/app/components/chat-section.tsx

@@ -1,57 +1,26 @@
 "use client";
 
+import { ChatSection as ChatSectionUI } from "@llamaindex/chat-ui";
+import "@llamaindex/chat-ui/styles/code.css";
+import "@llamaindex/chat-ui/styles/katex.css";
 import { useChat } from "ai/react";
-import { useState } from "react";
-import { ChatInput, ChatMessages } from "./ui/chat";
+import CustomChatInput from "./ui/chat/chat-input";
+import CustomChatMessages from "./ui/chat/chat-messages";
 import { useClientConfig } from "./ui/chat/hooks/use-config";
 
 export default function ChatSection() {
   const { backend } = useClientConfig();
-  const [requestData, setRequestData] = useState<any>();
-  const {
-    messages,
-    input,
-    isLoading,
-    handleSubmit,
-    handleInputChange,
-    reload,
-    stop,
-    append,
-    setInput,
-  } = useChat({
-    body: { data: requestData },
+  const handler = useChat({
     api: `${backend}/api/chat`,
-    headers: {
-      "Content-Type": "application/json", // using JSON because of vercel/ai 2.2.26
-    },
     onError: (error: unknown) => {
       if (!(error instanceof Error)) throw error;
-      const message = JSON.parse(error.message);
-      alert(message.detail);
+      alert(JSON.parse(error.message).detail);
     },
-    sendExtraMessageFields: true,
   });
-
   return (
-    <div className="space-y-4 w-full h-full flex flex-col">
-      <ChatMessages
-        messages={messages}
-        isLoading={isLoading}
-        reload={reload}
-        stop={stop}
-        append={append}
-      />
-      <ChatInput
-        input={input}
-        handleSubmit={handleSubmit}
-        handleInputChange={handleInputChange}
-        isLoading={isLoading}
-        messages={messages}
-        append={append}
-        setInput={setInput}
-        requestParams={{ params: requestData }}
-        setRequestData={setRequestData}
-      />
-    </div>
+    <ChatSectionUI handler={handler} className="w-full h-full">
+      <CustomChatMessages />
+      <CustomChatInput />
+    </ChatSectionUI>
   );
 }

templates/types/streaming/nextjs/app/components/ui/README-template.md

@@ -1 +0,0 @@
-Using the chat component from https://github.com/marcusschiesser/ui (based on https://ui.shadcn.com/)

templates/types/streaming/nextjs/app/components/ui/chat/chat-actions.tsx

@@ -1,28 +0,0 @@
-import { PauseCircle, RefreshCw } from "lucide-react";
-
-import { Button } from "../button";
-import { ChatHandler } from "./chat.interface";
-
-export default function ChatActions(
-  props: Pick<ChatHandler, "stop" | "reload"> & {
-    showReload?: boolean;
-    showStop?: boolean;
-  },
-) {
-  return (
-    <div className="space-x-4">
-      {props.showStop && (
-        <Button variant="outline" size="sm" onClick={props.stop}>
-          <PauseCircle className="mr-2 h-4 w-4" />
-          Stop generating
-        </Button>
-      )}
-      {props.showReload && (
-        <Button variant="outline" size="sm" onClick={props.reload}>
-          <RefreshCw className="mr-2 h-4 w-4" />
-          Regenerate
-        </Button>
-      )}
-    </div>
-  );
-}

templates/types/streaming/nextjs/app/components/ui/chat/chat-avatar.tsx

@@ -1,8 +1,10 @@
+import { useChatMessage } from "@llamaindex/chat-ui";
 import { User2 } from "lucide-react";
 import Image from "next/image";
 
-export default function ChatAvatar({ role }: { role: string }) {
-  if (role === "user") {
+export function ChatMessageAvatar() {
+  const { message } = useChatMessage();
+  if (message.role === "user") {
     return (
       <div className="flex h-8 w-8 shrink-0 select-none items-center justify-center rounded-md border bg-background shadow">
         <User2 className="h-4 w-4" />

templates/types/streaming/nextjs/app/components/ui/chat/chat-input.tsx

@@ -1,34 +1,13 @@
-import { JSONValue } from "ai";
-import React from "react";
-import { DocumentFile } from ".";
-import { Button } from "../button";
-import { DocumentPreview } from "../document-preview";
-import FileUploader from "../file-uploader";
-import { Textarea } from "../textarea";
-import UploadImagePreview from "../upload-image-preview";
-import { ChatHandler } from "./chat.interface";
-import { useFile } from "./hooks/use-file";
-import { LlamaCloudSelector } from "./widgets/LlamaCloudSelector";
+"use client";
 
-const ALLOWED_EXTENSIONS = ["png", "jpg", "jpeg", "csv", "pdf", "txt", "docx"];
+import { ChatInput, useChatUI, useFile } from "@llamaindex/chat-ui";
+import { DocumentPreview, ImagePreview } from "@llamaindex/chat-ui/widgets";
+import { LlamaCloudSelector } from "./custom/llama-cloud-selector";
+import { useClientConfig } from "./hooks/use-config";
 
-export default function ChatInput(
-  props: Pick<
-    ChatHandler,
-    | "isLoading"
-    | "input"
-    | "onFileUpload"
-    | "onFileError"
-    | "handleSubmit"
-    | "handleInputChange"
-    | "messages"
-    | "setInput"
-    | "append"
-  > & {
-    requestParams?: any;
-    setRequestData?: React.Dispatch<any>;
-  },
-) {
+export default function CustomChatInput() {
+  const { requestData, isLoading, input } = useChatUI();
+  const { backend } = useClientConfig();
   const {
     imageUrl,
     setImageUrl,
@@ -37,107 +16,65 @@ export default function ChatInput(
     removeDoc,
     reset,
     getAnnotations,
-  } = useFile();
-
-  // default submit function does not handle including annotations in the message
-  // so we need to use append function to submit new message with annotations
-  const handleSubmitWithAnnotations = (
-    e: React.FormEvent<HTMLFormElement>,
-    annotations: JSONValue[] | undefined,
-  ) => {
-    e.preventDefault();
-    props.append!({
-      content: props.input,
-      role: "user",
-      createdAt: new Date(),
-      annotations,
-    });
-    props.setInput!("");
-  };
-
-  const onSubmit = (e: React.FormEvent<HTMLFormElement>) => {
-    e.preventDefault();
-    const annotations = getAnnotations();
-    if (annotations.length) {
-      handleSubmitWithAnnotations(e, annotations);
-      return reset();
-    }
-    props.handleSubmit(e);
-  };
+  } = useFile({ uploadAPI: `${backend}/api/chat/upload` });
 
+  /**
+   * Handles file uploads. Overwrite to hook into the file upload behavior.
+   * @param file The file to upload
+   */
   const handleUploadFile = async (file: File) => {
+    // There's already an image uploaded, only allow one image at a time
     if (imageUrl) {
       alert("You can only upload one image at a time.");
       return;
     }
+
     try {
-      await uploadFile(file, props.requestParams);
-      props.onFileUpload?.(file);
+      // Upload the file and send with it the current request data
+      await uploadFile(file, requestData);
     } catch (error: any) {
-      const onFileUploadError = props.onFileError || window.alert;
-      onFileUploadError(error.message);
+      // Show error message if upload fails
+      alert(error.message);
     }
   };
 
-  const handleKeyDown = (e: React.KeyboardEvent<HTMLTextAreaElement>) => {
-    if (e.key === "Enter" && !e.shiftKey) {
-      e.preventDefault();
-      onSubmit(e as unknown as React.FormEvent<HTMLFormElement>);
-    }
-  };
+  // Get references to the upload files in message annotations format, see https://github.com/run-llama/chat-ui/blob/main/packages/chat-ui/src/hook/use-file.tsx#L56
+  const annotations = getAnnotations();
 
   return (
-    <form
-      onSubmit={onSubmit}
-      className="rounded-xl bg-white p-4 shadow-xl space-y-4 shrink-0"
+    <ChatInput
+      className="shadow-xl rounded-xl"
+      resetUploadedFiles={reset}
+      annotations={annotations}
     >
-      {imageUrl && (
-        <UploadImagePreview url={imageUrl} onRemove={() => setImageUrl(null)} />
-      )}
-      {files.length > 0 && (
-        <div className="flex gap-4 w-full overflow-auto py-2">
-          {files.map((file: DocumentFile) => (
-            <DocumentPreview
-              key={file.id}
-              file={file}
-              onRemove={() => removeDoc(file)}
-            />
-          ))}
-        </div>
-      )}
-      <div className="flex w-full items-start justify-between gap-4 ">
-        <Textarea
-          id="chat-input"
-          autoFocus
-          name="message"
-          placeholder="Type a message"
-          className="flex-1 min-h-0 h-[40px]"
-          value={props.input}
-          onChange={props.handleInputChange}
-          onKeyDown={handleKeyDown}
-        />
-        <FileUploader
-          onFileUpload={handleUploadFile}
-          onFileError={props.onFileError}
-          config={{
-            allowedExtensions: ALLOWED_EXTENSIONS,
-            disabled: props.isLoading,
-            multiple: true,
-          }}
-        />
-        {process.env.NEXT_PUBLIC_USE_LLAMACLOUD === "true" &&
-          props.setRequestData && (
-            <LlamaCloudSelector setRequestData={props.setRequestData} />
-          )}
-        <Button
-          type="submit"
-          disabled={
-            props.isLoading || (!props.input.trim() && files.length === 0)
-          }
-        >
-          Send message
-        </Button>
+      <div>
+        {/* Image preview section */}
+        {imageUrl && (
+          <ImagePreview url={imageUrl} onRemove={() => setImageUrl(null)} />
+        )}
+        {/* Document previews section */}
+        {files.length > 0 && (
+          <div className="flex gap-4 w-full overflow-auto py-2">
+            {files.map((file) => (
+              <DocumentPreview
+                key={file.id}
+                file={file}
+                onRemove={() => removeDoc(file)}
+              />
+            ))}
+          </div>
+        )}
       </div>
-    </form>
+      <ChatInput.Form>
+        <ChatInput.Field />
+        <ChatInput.Upload onUpload={handleUploadFile} />
+        <LlamaCloudSelector />
+        <ChatInput.Submit
+          disabled={
+            isLoading || (!input.trim() && files.length === 0 && !imageUrl)
+          }
+        />
+      </ChatInput.Form>
+    </ChatInput>
   );
 }

templates/types/streaming/nextjs/app/components/ui/chat/chat-message-content.tsx

@@ -0,0 +1,30 @@
+import {
+  ChatMessage,
+  ContentPosition,
+  getSourceAnnotationData,
+  useChatMessage,
+} from "@llamaindex/chat-ui";
+import { Markdown } from "./custom/markdown";
+import { ToolAnnotations } from "./tools/chat-tools";
+
+export function ChatMessageContent() {
+  const { message } = useChatMessage();
+  const customContent = [
+    {
+      // override the default markdown component
+      position: ContentPosition.MARKDOWN,
+      component: (
+        <Markdown
+          content={message.content}
+          sources={getSourceAnnotationData(message.annotations)?.[0]}
+        />
+      ),
+    },
+    {
+      // add the tool annotations after events
+      position: ContentPosition.AFTER_EVENTS,
+      component: <ToolAnnotations message={message} />,
+    },
+  ];
+  return <ChatMessage.Content content={customContent} />;
+}

templates/types/streaming/nextjs/app/components/ui/chat/chat-message/chat-agent-events.tsx

@@ -1,222 +0,0 @@
-import { icons, LucideIcon } from "lucide-react";
-import { useMemo } from "react";
-import { Button } from "../../button";
-import {
-  Drawer,
-  DrawerClose,
-  DrawerContent,
-  DrawerHeader,
-  DrawerTitle,
-  DrawerTrigger,
-} from "../../drawer";
-import { Progress } from "../../progress";
-import { AgentEventData, ProgressData } from "../index";
-import Markdown from "./markdown";
-
-const AgentIcons: Record<string, LucideIcon> = {
-  bot: icons.Bot,
-  researcher: icons.ScanSearch,
-  writer: icons.PenLine,
-  reviewer: icons.MessageCircle,
-  publisher: icons.BookCheck,
-};
-
-type StepText = {
-  text: string;
-};
-
-type StepProgress = {
-  text: string;
-  progress: ProgressData;
-};
-
-type MergedEvent = {
-  agent: string;
-  icon: LucideIcon;
-  steps: Array<StepText | StepProgress>;
-};
-
-export function ChatAgentEvents({
-  data,
-  isFinished,
-}: {
-  data: AgentEventData[];
-  isFinished: boolean;
-}) {
-  const events = useMemo(() => mergeAdjacentEvents(data), [data]);
-  return (
-    <div className="pl-2">
-      <div className="text-sm space-y-4">
-        {events.map((eventItem, index) => (
-          <AgentEventContent
-            key={index}
-            event={eventItem}
-            isLast={index === events.length - 1}
-            isFinished={isFinished}
-          />
-        ))}
-      </div>
-    </div>
-  );
-}
-
-const MAX_TEXT_LENGTH = 150;
-
-function TextContent({ agent, step }: { agent: string; step: StepText }) {
-  const { displayText, showMore } = useMemo(
-    () => ({
-      displayText: step.text.slice(0, MAX_TEXT_LENGTH),
-      showMore: step.text.length > MAX_TEXT_LENGTH,
-    }),
-    [step.text],
-  );
-
-  return (
-    <>
-      <div className="whitespace-break-spaces">
-        {!showMore && <span>{step.text}</span>}
-        {showMore && (
-          <div>
-            <span>{displayText}...</span>
-            <AgentEventDialog content={step.text} title={`Agent "${agent}"`}>
-              <span className="font-semibold underline cursor-pointer ml-2">
-                Show more
-              </span>
-            </AgentEventDialog>
-          </div>
-        )}
-      </div>
-    </>
-  );
-}
-
-function ProgressContent({ step }: { step: StepProgress }) {
-  const progressValue =
-    step.progress.total !== 0
-      ? Math.round(((step.progress.current + 1) / step.progress.total) * 100)
-      : 0;
-
-  return (
-    <div className="space-y-2 mt-2">
-      {step.text && (
-        <p className="text-sm text-muted-foreground">{step.text}</p>
-      )}
-      <Progress value={progressValue} className="w-full h-2" />
-      <p className="text-sm text-muted-foreground">
-        Processing {step.progress.current + 1} of {step.progress.total} steps...
-      </p>
-    </div>
-  );
-}
-
-function AgentEventContent({
-  event,
-  isLast,
-  isFinished,
-}: {
-  event: MergedEvent;
-  isLast: boolean;
-  isFinished: boolean;
-}) {
-  const { agent, steps } = event;
-  const AgentIcon = event.icon;
-  const textSteps = steps.filter((step) => !("progress" in step));
-  const progressSteps = steps.filter(
-    (step) => "progress" in step,
-  ) as StepProgress[];
-  // We only show progress at the last step
-  // TODO: once we support steps that work in parallel, we need to update this
-  const lastProgressStep =
-    progressSteps.length > 0
-      ? progressSteps[progressSteps.length - 1]
-      : undefined;
-
-  return (
-    <div className="flex gap-4 border-b pb-4 items-center fadein-agent">
-      <div className="w-[100px] flex flex-col items-center gap-2">
-        <div className="relative">
-          {isLast && !isFinished && (
-            <div className="absolute -top-0 -right-4">
-              <span className="relative flex h-3 w-3">
-                <span className="animate-ping absolute inline-flex h-full w-full rounded-full bg-sky-400 opacity-75"></span>
-                <span className="relative inline-flex rounded-full h-3 w-3 bg-sky-500"></span>
-              </span>
-            </div>
-          )}
-          <AgentIcon />
-        </div>
-        <span className="font-bold">{agent}</span>
-      </div>
-      {textSteps.length > 0 && (
-        <div className="flex-1">
-          <ul className="list-decimal space-y-2">
-            {textSteps.map((step, index) => (
-              <li key={index}>
-                <TextContent agent={agent} step={step} />
-              </li>
-            ))}
-          </ul>
-          {lastProgressStep && !isFinished && (
-            <ProgressContent step={lastProgressStep} />
-          )}
-        </div>
-      )}
-    </div>
-  );
-}
-
-type AgentEventDialogProps = {
-  title: string;
-  content: string;
-  children: React.ReactNode;
-};
-
-function AgentEventDialog(props: AgentEventDialogProps) {
-  return (
-    <Drawer direction="left">
-      <DrawerTrigger asChild>{props.children}</DrawerTrigger>
-      <DrawerContent className="w-3/5 mt-24 h-full max-h-[96%] ">
-        <DrawerHeader className="flex justify-between">
-          <div className="space-y-2">
-            <DrawerTitle>{props.title}</DrawerTitle>
-          </div>
-          <DrawerClose asChild>
-            <Button variant="outline">Close</Button>
-          </DrawerClose>
-        </DrawerHeader>
-        <div className="m-4 overflow-auto">
-          <Markdown content={props.content} />
-        </div>
-      </DrawerContent>
-    </Drawer>
-  );
-}
-
-function mergeAdjacentEvents(events: AgentEventData[]): MergedEvent[] {
-  const mergedEvents: MergedEvent[] = [];
-
-  for (const event of events) {
-    const lastMergedEvent = mergedEvents[mergedEvents.length - 1];
-
-    const eventStep: StepText | StepProgress = event.data
-      ? ({
-          text: event.text,
-          progress: event.data,
-        } as StepProgress)
-      : ({
-          text: event.text,
-        } as StepText);
-
-    if (lastMergedEvent && lastMergedEvent.agent === event.agent) {
-      lastMergedEvent.steps.push(eventStep);
-    } else {
-      mergedEvents.push({
-        agent: event.agent,
-        steps: [eventStep],
-        icon: AgentIcons[event.agent.toLowerCase()] ?? icons.Bot,
-      });
-    }
-  }
-
-  return mergedEvents;
-}

templates/types/streaming/nextjs/app/components/ui/chat/chat-message/chat-events.tsx

@@ -1,50 +0,0 @@
-import { ChevronDown, ChevronRight, Loader2 } from "lucide-react";
-import { useState } from "react";
-import { Button } from "../../button";
-import {
-  Collapsible,
-  CollapsibleContent,
-  CollapsibleTrigger,
-} from "../../collapsible";
-import { EventData } from "../index";
-
-export function ChatEvents({
-  data,
-  isLoading,
-}: {
-  data: EventData[];
-  isLoading: boolean;
-}) {
-  const [isOpen, setIsOpen] = useState(false);
-
-  const buttonLabel = isOpen ? "Hide events" : "Show events";
-
-  const EventIcon = isOpen ? (
-    <ChevronDown className="h-4 w-4" />
-  ) : (
-    <ChevronRight className="h-4 w-4" />
-  );
-
-  return (
-    <div className="border-l-2 border-indigo-400 pl-2">
-      <Collapsible open={isOpen} onOpenChange={setIsOpen}>
-        <CollapsibleTrigger asChild>
-          <Button variant="secondary" className="space-x-2">
-            {isLoading ? <Loader2 className="h-4 w-4 animate-spin" /> : null}
-            <span>{buttonLabel}</span>
-            {EventIcon}
-          </Button>
-        </CollapsibleTrigger>
-        <CollapsibleContent asChild>
-          <div className="mt-4 text-sm space-y-2">
-            {data.map((eventItem, index) => (
-              <div className="whitespace-break-spaces" key={index}>
-                {eventItem.title}
-              </div>
-            ))}
-          </div>
-        </CollapsibleContent>
-      </Collapsible>
-    </div>
-  );
-}

templates/types/streaming/nextjs/app/components/ui/chat/chat-message/chat-files.tsx

@@ -1,13 +0,0 @@
-import { DocumentPreview } from "../../document-preview";
-import { DocumentFileData } from "../index";
-
-export function ChatFiles({ data }: { data: DocumentFileData }) {
-  if (!data.files.length) return null;
-  return (
-    <div className="flex gap-2 items-center">
-      {data.files.map((file, index) => (
-        <DocumentPreview key={file.id} file={file} />
-      ))}
-    </div>
-  );
-}

templates/types/streaming/nextjs/app/components/ui/chat/chat-message/chat-image.tsx

@@ -1,17 +0,0 @@
-import Image from "next/image";
-import { type ImageData } from "../index";
-
-export function ChatImage({ data }: { data: ImageData }) {
-  return (
-    <div className="rounded-md max-w-[200px] shadow-md">
-      <Image
-        src={data.url}
-        width={0}
-        height={0}
-        sizes="100vw"
-        style={{ width: "100%", height: "auto" }}
-        alt=""
-      />
-    </div>
-  );
-}

templates/types/streaming/nextjs/app/components/ui/chat/chat-message/chat-sources.tsx

@@ -1,173 +0,0 @@
-import { Check, Copy } from "lucide-react";
-import { useMemo } from "react";
-import { Button } from "../../button";
-import { PreviewCard } from "../../document-preview";
-import {
-  HoverCard,
-  HoverCardContent,
-  HoverCardTrigger,
-} from "../../hover-card";
-import { cn } from "../../lib/utils";
-import { useCopyToClipboard } from "../hooks/use-copy-to-clipboard";
-import { DocumentFileType, SourceData, SourceNode } from "../index";
-import PdfDialog from "../widgets/PdfDialog";
-
-type Document = {
-  url: string;
-  sources: SourceNode[];
-};
-
-export function ChatSources({ data }: { data: SourceData }) {
-  const documents: Document[] = useMemo(() => {
-    // group nodes by document (a document must have a URL)
-    const nodesByUrl: Record<string, SourceNode[]> = {};
-    data.nodes.forEach((node) => {
-      const key = node.url;
-      nodesByUrl[key] ??= [];
-      nodesByUrl[key].push(node);
-    });
-
-    // convert to array of documents
-    return Object.entries(nodesByUrl).map(([url, sources]) => ({
-      url,
-      sources,
-    }));
-  }, [data.nodes]);
-
-  if (documents.length === 0) return null;
-
-  return (
-    <div className="space-y-2 text-sm">
-      <div className="font-semibold text-lg">Sources:</div>
-      <div className="flex gap-3 flex-wrap">
-        {documents.map((document) => {
-          return <DocumentInfo key={document.url} document={document} />;
-        })}
-      </div>
-    </div>
-  );
-}
-
-function SourceInfo({ node, index }: { node?: SourceNode; index: number }) {
-  if (!node) return <SourceNumberButton index={index} />;
-  return (
-    <HoverCard>
-      <HoverCardTrigger
-        className="cursor-default"
-        onClick={(e) => {
-          e.preventDefault();
-          e.stopPropagation();
-        }}
-      >
-        <SourceNumberButton
-          index={index}
-          className="hover:text-white hover:bg-primary"
-        />
-      </HoverCardTrigger>
-      <HoverCardContent className="w-[400px]">
-        <NodeInfo nodeInfo={node} />
-      </HoverCardContent>
-    </HoverCard>
-  );
-}
-
-export function SourceNumberButton({
-  index,
-  className,
-}: {
-  index: number;
-  className?: string;
-}) {
-  return (
-    <span
-      className={cn(
-        "text-xs w-5 h-5 rounded-full bg-gray-100 inline-flex items-center justify-center",
-        className,
-      )}
-    >
-      {index + 1}
-    </span>
-  );
-}
-
-export function DocumentInfo({
-  document,
-  className,
-}: {
-  document: Document;
-  className?: string;
-}) {
-  const { url, sources } = document;
-  // Extract filename from URL
-  const urlParts = url.split("/");
-  const fileName = urlParts.length > 0 ? urlParts[urlParts.length - 1] : url;
-  const fileExt = fileName?.split(".").pop() as DocumentFileType | undefined;
-
-  const previewFile = {
-    name: fileName,
-    type: fileExt as DocumentFileType,
-  };
-
-  const DocumentDetail = (
-    <div className={`relative ${className}`}>
-      <PreviewCard className={"cursor-pointer"} file={previewFile} />
-      <div className="absolute bottom-2 right-2 space-x-2 flex">
-        {sources.map((node: SourceNode, index: number) => (
-          <div key={node.id}>
-            <SourceInfo node={node} index={index} />
-          </div>
-        ))}
-      </div>
-    </div>
-  );
-
-  if (url.endsWith(".pdf")) {
-    // open internal pdf dialog for pdf files when click document card
-    return <PdfDialog documentId={url} url={url} trigger={DocumentDetail} />;
-  }
-  // open external link when click document card for other file types
-  return <div onClick={() => window.open(url, "_blank")}>{DocumentDetail}</div>;
-}
-
-function NodeInfo({ nodeInfo }: { nodeInfo: SourceNode }) {
-  const { isCopied, copyToClipboard } = useCopyToClipboard({ timeout: 1000 });
-
-  const pageNumber =
-    // XXX: page_label is used in Python, but page_number is used by Typescript
-    (nodeInfo.metadata?.page_number as number) ??
-    (nodeInfo.metadata?.page_label as number) ??
-    null;
-
-  return (
-    <div className="space-y-4">
-      <div className="flex justify-between items-center">
-        <span className="font-semibold">
-          {pageNumber ? `On page ${pageNumber}:` : "Node content:"}
-        </span>
-        {nodeInfo.text && (
-          <Button
-            onClick={(e) => {
-              e.stopPropagation();
-              copyToClipboard(nodeInfo.text);
-            }}
-            size="icon"
-            variant="ghost"
-            className="h-12 w-12 shrink-0"
-          >
-            {isCopied ? (
-              <Check className="h-4 w-4" />
-            ) : (
-              <Copy className="h-4 w-4" />
-            )}
-          </Button>
-        )}
-      </div>
-
-      {nodeInfo.text && (
-        <pre className="max-h-[200px] overflow-auto whitespace-pre-line">
-          &ldquo;{nodeInfo.text}&rdquo;
-        </pre>
-      )}
-    </div>
-  );
-}

templates/types/streaming/nextjs/app/components/ui/chat/chat-message/chat-suggestedQuestions.tsx

@@ -1,31 +0,0 @@
-import { ChatHandler, SuggestedQuestionsData } from "..";
-
-export function SuggestedQuestions({
-  questions,
-  append,
-  isLastMessage,
-}: {
-  questions: SuggestedQuestionsData;
-  append: Pick<ChatHandler, "append">["append"];
-  isLastMessage: boolean;
-}) {
-  const showQuestions = isLastMessage && questions.length > 0;
-  return (
-    showQuestions &&
-    append !== undefined && (
-      <div className="flex flex-col space-y-2">
-        {questions.map((question, index) => (
-          <a
-            key={index}
-            onClick={() => {
-              append({ role: "user", content: question });
-            }}
-            className="text-sm italic hover:underline cursor-pointer"
-          >
-            {"->"} {question}
-          </a>
-        ))}
-      </div>
-    )
-  );
-}

templates/types/streaming/nextjs/app/components/ui/chat/chat-message/chat-tools.tsx

@@ -1,40 +0,0 @@
-import { ToolData } from "../index";
-import { Artifact, CodeArtifact } from "../widgets/Artifact";
-import { WeatherCard, WeatherData } from "../widgets/WeatherCard";
-
-// TODO: If needed, add displaying more tool outputs here
-export default function ChatTools({
-  data,
-  artifactVersion,
-}: {
-  data: ToolData;
-  artifactVersion?: number;
-}) {
-  if (!data) return null;
-  const { toolCall, toolOutput } = data;
-
-  if (toolOutput.isError) {
-    return (
-      <div className="border-l-2 border-red-400 pl-2">
-        There was an error when calling the tool {toolCall.name} with input:{" "}
-        <br />
-        {JSON.stringify(toolCall.input)}
-      </div>
-    );
-  }
-
-  switch (toolCall.name) {
-    case "get_weather_information":
-      const weatherData = toolOutput.output as unknown as WeatherData;
-      return <WeatherCard data={weatherData} />;
-    case "artifact":
-      return (
-        <Artifact
-          artifact={toolOutput.output as CodeArtifact}
-          version={artifactVersion}
-        />
-      );
-    default:
-      return null;
-  }
-}

templates/types/streaming/nextjs/app/components/ui/chat/chat-message/codeblock.tsx

@@ -1,131 +0,0 @@
-"use client";
-
-import hljs from "highlight.js";
-// instead of atom-one-dark theme, there are a lot of others: https://highlightjs.org/demo
-import "highlight.js/styles/atom-one-dark-reasonable.css";
-import { Check, Copy, Download } from "lucide-react";
-import { FC, memo, useEffect, useRef } from "react";
-import { Button } from "../../button";
-import { useCopyToClipboard } from "../hooks/use-copy-to-clipboard";
-
-interface Props {
-  language: string;
-  value: string;
-  className?: string;
-}
-
-interface languageMap {
-  [key: string]: string | undefined;
-}
-
-export const programmingLanguages: languageMap = {
-  javascript: ".js",
-  python: ".py",
-  java: ".java",
-  c: ".c",
-  cpp: ".cpp",
-  "c++": ".cpp",
-  "c#": ".cs",
-  ruby: ".rb",
-  php: ".php",
-  swift: ".swift",
-  "objective-c": ".m",
-  kotlin: ".kt",
-  typescript: ".ts",
-  go: ".go",
-  perl: ".pl",
-  rust: ".rs",
-  scala: ".scala",
-  haskell: ".hs",
-  lua: ".lua",
-  shell: ".sh",
-  sql: ".sql",
-  html: ".html",
-  css: ".css",
-  // add more file extensions here, make sure the key is same as language prop in CodeBlock.tsx component
-};
-
-export const generateRandomString = (length: number, lowercase = false) => {
-  const chars = "ABCDEFGHJKLMNPQRSTUVWXY3456789"; // excluding similar looking characters like Z, 2, I, 1, O, 0
-  let result = "";
-  for (let i = 0; i < length; i++) {
-    result += chars.charAt(Math.floor(Math.random() * chars.length));
-  }
-  return lowercase ? result.toLowerCase() : result;
-};
-
-const CodeBlock: FC<Props> = memo(({ language, value, className }) => {
-  const { isCopied, copyToClipboard } = useCopyToClipboard({ timeout: 2000 });
-  const codeRef = useRef<HTMLElement>(null);
-
-  useEffect(() => {
-    if (codeRef.current && codeRef.current.dataset.highlighted !== "yes") {
-      hljs.highlightElement(codeRef.current);
-    }
-  }, [language, value]);
-
-  const downloadAsFile = () => {
-    if (typeof window === "undefined") {
-      return;
-    }
-    const fileExtension = programmingLanguages[language] || ".file";
-    const suggestedFileName = `file-${generateRandomString(
-      3,
-      true,
-    )}${fileExtension}`;
-    const fileName = window.prompt("Enter file name", suggestedFileName);
-
-    if (!fileName) {
-      // User pressed cancel on prompt.
-      return;
-    }
-
-    const blob = new Blob([value], { type: "text/plain" });
-    const url = URL.createObjectURL(blob);
-    const link = document.createElement("a");
-    link.download = fileName;
-    link.href = url;
-    link.style.display = "none";
-    document.body.appendChild(link);
-    link.click();
-    document.body.removeChild(link);
-    URL.revokeObjectURL(url);
-  };
-
-  const onCopy = () => {
-    if (isCopied) return;
-    copyToClipboard(value);
-  };
-
-  return (
-    <div
-      className={`codeblock relative w-full bg-zinc-950 font-sans ${className}`}
-    >
-      <div className="flex w-full items-center justify-between bg-zinc-800 px-6 py-2 pr-4 text-zinc-100">
-        <span className="text-xs lowercase">{language}</span>
-        <div className="flex items-center space-x-1">
-          <Button variant="ghost" onClick={downloadAsFile} size="icon">
-            <Download />
-            <span className="sr-only">Download</span>
-          </Button>
-          <Button variant="ghost" size="icon" onClick={onCopy}>
-            {isCopied ? (
-              <Check className="h-4 w-4" />
-            ) : (
-              <Copy className="h-4 w-4" />
-            )}
-            <span className="sr-only">Copy code</span>
-          </Button>
-        </div>
-      </div>
-      <pre className="border border-zinc-700">
-        <code ref={codeRef} className={`language-${language} font-mono`}>
-          {value}
-        </code>
-      </pre>
-    </div>
-  );
-});
-CodeBlock.displayName = "CodeBlock";
-
-export { CodeBlock };

templates/types/streaming/nextjs/app/components/ui/chat/chat-message/index.tsx

@@ -1,184 +0,0 @@
-import { Check, Copy } from "lucide-react";
-
-import { Message } from "ai";
-import { Fragment } from "react";
-import { Button } from "../../button";
-import { useCopyToClipboard } from "../hooks/use-copy-to-clipboard";
-import {
-  AgentEventData,
-  ChatHandler,
-  DocumentFileData,
-  EventData,
-  ImageData,
-  MessageAnnotation,
-  MessageAnnotationType,
-  SuggestedQuestionsData,
-  ToolData,
-  getAnnotationData,
-  getSourceAnnotationData,
-} from "../index";
-import { ChatAgentEvents } from "./chat-agent-events";
-import ChatAvatar from "./chat-avatar";
-import { ChatEvents } from "./chat-events";
-import { ChatFiles } from "./chat-files";
-import { ChatImage } from "./chat-image";
-import { ChatSources } from "./chat-sources";
-import { SuggestedQuestions } from "./chat-suggestedQuestions";
-import ChatTools from "./chat-tools";
-import Markdown from "./markdown";
-
-type ContentDisplayConfig = {
-  order: number;
-  component: JSX.Element | null;
-};
-
-function ChatMessageContent({
-  message,
-  isLoading,
-  append,
-  isLastMessage,
-  artifactVersion,
-}: {
-  message: Message;
-  isLoading: boolean;
-  append: Pick<ChatHandler, "append">["append"];
-  isLastMessage: boolean;
-  artifactVersion: number | undefined;
-}) {
-  const annotations = message.annotations as MessageAnnotation[] | undefined;
-  if (!annotations?.length) return <Markdown content={message.content} />;
-
-  const imageData = getAnnotationData<ImageData>(
-    annotations,
-    MessageAnnotationType.IMAGE,
-  );
-  const contentFileData = getAnnotationData<DocumentFileData>(
-    annotations,
-    MessageAnnotationType.DOCUMENT_FILE,
-  );
-  const eventData = getAnnotationData<EventData>(
-    annotations,
-    MessageAnnotationType.EVENTS,
-  );
-  const agentEventData = getAnnotationData<AgentEventData>(
-    annotations,
-    MessageAnnotationType.AGENT_EVENTS,
-  );
-
-  const sourceData = getSourceAnnotationData(annotations);
-
-  const toolData = getAnnotationData<ToolData>(
-    annotations,
-    MessageAnnotationType.TOOLS,
-  );
-  const suggestedQuestionsData = getAnnotationData<SuggestedQuestionsData>(
-    annotations,
-    MessageAnnotationType.SUGGESTED_QUESTIONS,
-  );
-
-  const contents: ContentDisplayConfig[] = [
-    {
-      order: 1,
-      component: imageData[0] ? <ChatImage data={imageData[0]} /> : null,
-    },
-    {
-      order: -3,
-      component:
-        eventData.length > 0 ? (
-          <ChatEvents isLoading={isLoading} data={eventData} />
-        ) : null,
-    },
-    {
-      order: -2,
-      component:
-        agentEventData.length > 0 ? (
-          <ChatAgentEvents
-            data={agentEventData}
-            isFinished={!!message.content}
-          />
-        ) : null,
-    },
-    {
-      order: 2,
-      component: contentFileData[0] ? (
-        <ChatFiles data={contentFileData[0]} />
-      ) : null,
-    },
-    {
-      order: -1,
-      component: toolData[0] ? (
-        <ChatTools data={toolData[0]} artifactVersion={artifactVersion} />
-      ) : null,
-    },
-    {
-      order: 0,
-      component: <Markdown content={message.content} sources={sourceData[0]} />,
-    },
-    {
-      order: 3,
-      component: sourceData[0] ? <ChatSources data={sourceData[0]} /> : null,
-    },
-    {
-      order: 4,
-      component: suggestedQuestionsData[0] ? (
-        <SuggestedQuestions
-          questions={suggestedQuestionsData[0]}
-          append={append}
-          isLastMessage={isLastMessage}
-        />
-      ) : null,
-    },
-  ];
-
-  return (
-    <div className="flex-1 gap-4 flex flex-col">
-      {contents
-        .sort((a, b) => a.order - b.order)
-        .map((content, index) => (
-          <Fragment key={index}>{content.component}</Fragment>
-        ))}
-    </div>
-  );
-}
-
-export default function ChatMessage({
-  chatMessage,
-  isLoading,
-  append,
-  isLastMessage,
-  artifactVersion,
-}: {
-  chatMessage: Message;
-  isLoading: boolean;
-  append: Pick<ChatHandler, "append">["append"];
-  isLastMessage: boolean;
-  artifactVersion: number | undefined;
-}) {
-  const { isCopied, copyToClipboard } = useCopyToClipboard({ timeout: 2000 });
-  return (
-    <div className="flex items-start gap-4 pr-5 pt-5">
-      <ChatAvatar role={chatMessage.role} />
-      <div className="group flex flex-1 justify-between gap-2">
-        <ChatMessageContent
-          message={chatMessage}
-          isLoading={isLoading}
-          append={append}
-          isLastMessage={isLastMessage}
-          artifactVersion={artifactVersion}
-        />
-        <Button
-          onClick={() => copyToClipboard(chatMessage.content)}
-          size="icon"
-          variant="ghost"
-          className="h-8 w-8 opacity-0 group-hover:opacity-100"
-        >
-          {isCopied ? (
-            <Check className="h-4 w-4" />
-          ) : (
-            <Copy className="h-4 w-4" />
-          )}
-        </Button>
-      </div>
-    </div>
-  );
-}

templates/types/streaming/nextjs/app/components/ui/chat/chat-message/markdown.tsx

@@ -1,172 +0,0 @@
-import "katex/dist/katex.min.css";
-import { FC, memo } from "react";
-import ReactMarkdown, { Options } from "react-markdown";
-import rehypeKatex from "rehype-katex";
-import remarkGfm from "remark-gfm";
-import remarkMath from "remark-math";
-
-import { DOCUMENT_FILE_TYPES, DocumentFileType, SourceData } from "..";
-import { useClientConfig } from "../hooks/use-config";
-import { DocumentInfo, SourceNumberButton } from "./chat-sources";
-import { CodeBlock } from "./codeblock";
-
-const MemoizedReactMarkdown: FC<Options> = memo(
-  ReactMarkdown,
-  (prevProps, nextProps) =>
-    prevProps.children === nextProps.children &&
-    prevProps.className === nextProps.className,
-);
-
-const preprocessLaTeX = (content: string) => {
-  // Replace block-level LaTeX delimiters \[ \] with $$ $$
-  const blockProcessedContent = content.replace(
-    /\\\[([\s\S]*?)\\\]/g,
-    (_, equation) => `$$${equation}$$`,
-  );
-  // Replace inline LaTeX delimiters \( \) with $ $
-  const inlineProcessedContent = blockProcessedContent.replace(
-    /\\\[([\s\S]*?)\\\]/g,
-    (_, equation) => `$${equation}$`,
-  );
-  return inlineProcessedContent;
-};
-
-const preprocessMedia = (content: string) => {
-  // Remove `sandbox:` from the beginning of the URL
-  // to fix OpenAI's models issue appending `sandbox:` to the relative URL
-  return content.replace(/(sandbox|attachment|snt):/g, "");
-};
-
-/**
- * Update the citation flag [citation:id]() to the new format [citation:index](url)
- */
-const preprocessCitations = (content: string, sources?: SourceData) => {
-  if (sources) {
-    const citationRegex = /\[citation:(.+?)\]\(\)/g;
-    let match;
-    // Find all the citation references in the content
-    while ((match = citationRegex.exec(content)) !== null) {
-      const citationId = match[1];
-      // Find the source node with the id equal to the citation-id, also get the index of the source node
-      const sourceNode = sources.nodes.find((node) => node.id === citationId);
-      // If the source node is found, replace the citation reference with the new format
-      if (sourceNode !== undefined) {
-        content = content.replace(
-          match[0],
-          `[citation:${sources.nodes.indexOf(sourceNode)}]()`,
-        );
-      } else {
-        // If the source node is not found, remove the citation reference
-        content = content.replace(match[0], "");
-      }
-    }
-  }
-  return content;
-};
-
-const preprocessContent = (content: string, sources?: SourceData) => {
-  return preprocessCitations(
-    preprocessMedia(preprocessLaTeX(content)),
-    sources,
-  );
-};
-
-export default function Markdown({
-  content,
-  sources,
-}: {
-  content: string;
-  sources?: SourceData;
-}) {
-  const processedContent = preprocessContent(content, sources);
-  const { backend } = useClientConfig();
-
-  return (
-    <MemoizedReactMarkdown
-      className="prose dark:prose-invert prose-p:leading-relaxed prose-pre:p-0 break-words custom-markdown"
-      remarkPlugins={[remarkGfm, remarkMath]}
-      rehypePlugins={[rehypeKatex as any]}
-      components={{
-        p({ children }) {
-          return <div className="mb-2 last:mb-0">{children}</div>;
-        },
-        code({ node, inline, className, children, ...props }) {
-          if (children.length) {
-            if (children[0] == "▍") {
-              return (
-                <span className="mt-1 animate-pulse cursor-default">▍</span>
-              );
-            }
-
-            children[0] = (children[0] as string).replace("`▍`", "▍");
-          }
-
-          const match = /language-(\w+)/.exec(className || "");
-
-          if (inline) {
-            return (
-              <code className={className} {...props}>
-                {children}
-              </code>
-            );
-          }
-
-          return (
-            <CodeBlock
-              key={Math.random()}
-              language={(match && match[1]) || ""}
-              value={String(children).replace(/\n$/, "")}
-              className="mb-2"
-              {...props}
-            />
-          );
-        },
-        a({ href, children }) {
-          // If href starts with `{backend}/api/files`, then it's a local document and we use DocumenInfo for rendering
-          if (href?.startsWith(backend + "/api/files")) {
-            // Check if the file is document file type
-            const fileExtension = href.split(".").pop()?.toLowerCase();
-
-            if (
-              fileExtension &&
-              DOCUMENT_FILE_TYPES.includes(fileExtension as DocumentFileType)
-            ) {
-              return (
-                <DocumentInfo
-                  document={{
-                    url: backend
-                      ? new URL(decodeURIComponent(href)).href
-                      : href,
-                    sources: [],
-                  }}
-                  className="mb-2 mt-2"
-                />
-              );
-            }
-          }
-          // If a text link starts with 'citation:', then render it as a citation reference
-          if (
-            Array.isArray(children) &&
-            typeof children[0] === "string" &&
-            children[0].startsWith("citation:")
-          ) {
-            const index = Number(children[0].replace("citation:", ""));
-            if (!isNaN(index)) {
-              return <SourceNumberButton index={index} />;
-            } else {
-              // citation is not looked up yet, don't render anything
-              return <></>;
-            }
-          }
-          return (
-            <a href={href} target="_blank">
-              {children}
-            </a>
-          );
-        },
-      }}
-    >
-      {processedContent}
-    </MemoizedReactMarkdown>
-  );
-}

templates/types/streaming/nextjs/app/components/ui/chat/chat-messages.tsx

@@ -1,136 +1,30 @@
-import { Loader2 } from "lucide-react";
-import { useEffect, useMemo, useRef, useState } from "react";
+"use client";
 
-import { ToolData } from ".";
-import { Button } from "../button";
-import ChatActions from "./chat-actions";
-import ChatMessage from "./chat-message";
-import { ChatHandler } from "./chat.interface";
-import { useClientConfig } from "./hooks/use-config";
-
-export default function ChatMessages(
-  props: Pick<
-    ChatHandler,
-    "messages" | "isLoading" | "reload" | "stop" | "append"
-  >,
-) {
-  const { backend } = useClientConfig();
-  const [starterQuestions, setStarterQuestions] = useState<string[]>();
-
-  const scrollableChatContainerRef = useRef<HTMLDivElement>(null);
-  const messageLength = props.messages.length;
-  const lastMessage = props.messages[messageLength - 1];
-
-  const scrollToBottom = () => {
-    if (scrollableChatContainerRef.current) {
-      scrollableChatContainerRef.current.scrollTop =
-        scrollableChatContainerRef.current.scrollHeight;
-    }
-  };
-
-  const isLastMessageFromAssistant =
-    messageLength > 0 && lastMessage?.role !== "user";
-  const showReload =
-    props.reload && !props.isLoading && isLastMessageFromAssistant;
-  const showStop = props.stop && props.isLoading;
-
-  // `isPending` indicate
-  // that stream response is not yet received from the server,
-  // so we show a loading indicator to give a better UX.
-  const isPending = props.isLoading && !isLastMessageFromAssistant;
-
-  useEffect(() => {
-    scrollToBottom();
-  }, [messageLength, lastMessage]);
-
-  useEffect(() => {
-    if (!starterQuestions) {
-      fetch(`${backend}/api/chat/config`)
-        .then((response) => response.json())
-        .then((data) => {
-          if (data?.starterQuestions) {
-            setStarterQuestions(data.starterQuestions);
-          }
-        })
-        .catch((error) => console.error("Error fetching config", error));
-    }
-  }, [starterQuestions, backend]);
-
-  // build a map of message id to artifact version
-  const artifactVersionMap = useMemo(() => {
-    const map = new Map<string, number | undefined>();
-    let versionIndex = 1;
-    props.messages.forEach((m) => {
-      m.annotations?.forEach((annotation) => {
-        if (
-          typeof annotation === "object" &&
-          annotation != null &&
-          "type" in annotation &&
-          annotation.type === "tools"
-        ) {
-          const data = annotation.data as ToolData;
-          if (data?.toolCall?.name === "artifact") {
-            map.set(m.id, versionIndex);
-            versionIndex++;
-          }
-        }
-      });
-    });
-    return map;
-  }, [props.messages]);
+import { ChatMessage, ChatMessages, useChatUI } from "@llamaindex/chat-ui";
+import { ChatMessageAvatar } from "./chat-avatar";
+import { ChatMessageContent } from "./chat-message-content";
+import { ChatStarter } from "./chat-starter";
 
+export default function CustomChatMessages() {
+  const { messages } = useChatUI();
   return (
-    <div
-      className="flex-1 w-full rounded-xl bg-white p-4 shadow-xl relative overflow-y-auto"
-      ref={scrollableChatContainerRef}
-    >
-      <div className="flex flex-col gap-5 divide-y">
-        {props.messages.map((m, i) => {
-          const isLoadingMessage = i === messageLength - 1 && props.isLoading;
-          return (
-            <ChatMessage
-              key={m.id}
-              chatMessage={m}
-              isLoading={isLoadingMessage}
-              append={props.append!}
-              isLastMessage={i === messageLength - 1}
-              artifactVersion={artifactVersionMap.get(m.id)}
-            />
-          );
-        })}
-        {isPending && (
-          <div className="flex justify-center items-center pt-10">
-            <Loader2 className="h-4 w-4 animate-spin" />
-          </div>
-        )}
-      </div>
-      {(showReload || showStop) && (
-        <div className="flex justify-end py-4">
-          <ChatActions
-            reload={props.reload}
-            stop={props.stop}
-            showReload={showReload}
-            showStop={showStop}
-          />
-        </div>
-      )}
-      {!messageLength && starterQuestions?.length && props.append && (
-        <div className="absolute bottom-6 left-0 w-full">
-          <div className="grid grid-cols-2 gap-2 mx-20">
-            {starterQuestions.map((question, i) => (
-              <Button
-                variant="outline"
-                key={i}
-                onClick={() =>
-                  props.append!({ role: "user", content: question })
-                }
-              >
-                {question}
-              </Button>
-            ))}
-          </div>
-        </div>
-      )}
-    </div>
+    <ChatMessages className="shadow-xl rounded-xl">
+      <ChatMessages.List>
+        {messages.map((message, index) => (
+          <ChatMessage
+            key={message.id}
+            message={message}
+            isLast={index === messages.length - 1}
+          >
+            <ChatMessageAvatar />
+            <ChatMessageContent />
+            <ChatMessage.Actions />
+          </ChatMessage>
+        ))}
+        <ChatMessages.Loading />
+      </ChatMessages.List>
+      <ChatMessages.Actions />
+      <ChatStarter />
+    </ChatMessages>
   );
 }

templates/types/streaming/nextjs/app/components/ui/chat/chat-starter.tsx

@@ -0,0 +1,26 @@
+import { useChatUI } from "@llamaindex/chat-ui";
+import { StarterQuestions } from "@llamaindex/chat-ui/widgets";
+import { useEffect, useState } from "react";
+import { useClientConfig } from "./hooks/use-config";
+
+export function ChatStarter() {
+  const { append } = useChatUI();
+  const { backend } = useClientConfig();
+  const [starterQuestions, setStarterQuestions] = useState<string[]>();
+
+  useEffect(() => {
+    if (!starterQuestions) {
+      fetch(`${backend}/api/chat/config`)
+        .then((response) => response.json())
+        .then((data) => {
+          if (data?.starterQuestions) {
+            setStarterQuestions(data.starterQuestions);
+          }
+        })
+        .catch((error) => console.error("Error fetching config", error));
+    }
+  }, [starterQuestions, backend]);
+
+  if (!starterQuestions?.length) return null;
+  return <StarterQuestions append={append} questions={starterQuestions} />;
+}

templates/types/streaming/nextjs/app/components/ui/chat/chat.interface.ts

@@ -1,25 +0,0 @@
-import { Message } from "ai";
-
-export interface ChatHandler {
-  messages: Message[];
-  input: string;
-  isLoading: boolean;
-  handleSubmit: (
-    e: React.FormEvent<HTMLFormElement>,
-    ops?: {
-      data?: any;
-    },
-  ) => void;
-  handleInputChange: (e: React.ChangeEvent<HTMLTextAreaElement>) => void;
-  reload?: () => void;
-  stop?: () => void;
-  onFileUpload?: (file: File) => Promise<void>;
-  onFileError?: (errMsg: string) => void;
-  setInput?: (input: string) => void;
-  append?: (
-    message: Message | Omit<Message, "id">,
-    ops?: {
-      data: any;
-    },
-  ) => Promise<string | null | undefined>;
-}

templates/types/streaming/nextjs/app/components/ui/chat/custom/llama-cloud-selector.tsx

@@ -1,3 +1,4 @@
+import { useChatUI } from "@llamaindex/chat-ui";
 import { Loader2 } from "lucide-react";
 import { useCallback, useEffect, useState } from "react";
 import {
@@ -35,19 +36,18 @@ type LlamaCloudConfig = {
 };
 
 export interface LlamaCloudSelectorProps {
-  setRequestData?: React.Dispatch<any>;
   onSelect?: (pipeline: PipelineConfig | undefined) => void;
   defaultPipeline?: PipelineConfig;
   shouldCheckValid?: boolean;
 }
 
 export function LlamaCloudSelector({
-  setRequestData,
   onSelect,
   defaultPipeline,
   shouldCheckValid = false,
 }: LlamaCloudSelectorProps) {
   const { backend } = useClientConfig();
+  const { setRequestData } = useChatUI();
   const [config, setConfig] = useState<LlamaCloudConfig>();
 
   const updateRequestParams = useCallback(
@@ -97,6 +97,10 @@ export function LlamaCloudSelector({
     setPipeline(JSON.parse(value) as PipelineConfig);
   };
 
+  if (process.env.NEXT_PUBLIC_USE_LLAMACLOUD !== "true") {
+    return null;
+  }
+
   if (!config) {
     return (
       <div className="flex justify-center items-center p-3">

templates/types/streaming/nextjs/app/components/ui/chat/custom/markdown.tsx

@@ -0,0 +1,27 @@
+import { SourceData } from "@llamaindex/chat-ui";
+import { Markdown as MarkdownUI } from "@llamaindex/chat-ui/widgets";
+import { useClientConfig } from "../hooks/use-config";
+
+const preprocessMedia = (content: string) => {
+  // Remove `sandbox:` from the beginning of the URL before rendering markdown
+  // OpenAI models sometimes prepend `sandbox:` to relative URLs - this fixes it
+  return content.replace(/(sandbox|attachment|snt):/g, "");
+};
+
+export function Markdown({
+  content,
+  sources,
+}: {
+  content: string;
+  sources?: SourceData;
+}) {
+  const { backend } = useClientConfig();
+  const processedContent = preprocessMedia(content);
+  return (
+    <MarkdownUI
+      content={processedContent}
+      backend={backend}
+      sources={sources}
+    />
+  );
+}

templates/types/streaming/nextjs/app/components/ui/chat/hooks/use-file.ts

@@ -1,121 +0,0 @@
-"use client";
-
-import { JSONValue } from "llamaindex";
-import { useState } from "react";
-import {
-  DocumentFile,
-  DocumentFileType,
-  MessageAnnotation,
-  MessageAnnotationType,
-} from "..";
-import { useClientConfig } from "./use-config";
-
-const docMineTypeMap: Record<string, DocumentFileType> = {
-  "text/csv": "csv",
-  "application/pdf": "pdf",
-  "text/plain": "txt",
-  "application/vnd.openxmlformats-officedocument.wordprocessingml.document":
-    "docx",
-};
-
-export function useFile() {
-  const { backend } = useClientConfig();
-  const [imageUrl, setImageUrl] = useState<string | null>(null);
-  const [files, setFiles] = useState<DocumentFile[]>([]);
-
-  const addDoc = (file: DocumentFile) => {
-    const existedFile = files.find((f) => f.id === file.id);
-    if (!existedFile) {
-      setFiles((prev) => [...prev, file]);
-      return true;
-    }
-    return false;
-  };
-
-  const removeDoc = (file: DocumentFile) => {
-    setFiles((prev) => prev.filter((f) => f.id !== file.id));
-  };
-
-  const reset = () => {
-    imageUrl && setImageUrl(null);
-    files.length && setFiles([]);
-  };
-
-  const uploadContent = async (
-    file: File,
-    requestParams: any = {},
-  ): Promise<DocumentFile> => {
-    const base64 = await readContent({ file, asUrl: true });
-    const uploadAPI = `${backend}/api/chat/upload`;
-    const response = await fetch(uploadAPI, {
-      method: "POST",
-      headers: {
-        "Content-Type": "application/json",
-      },
-      body: JSON.stringify({
-        ...requestParams,
-        base64,
-        name: file.name,
-      }),
-    });
-    if (!response.ok) throw new Error("Failed to upload document.");
-    return (await response.json()) as DocumentFile;
-  };
-
-  const getAnnotations = () => {
-    const annotations: MessageAnnotation[] = [];
-    if (imageUrl) {
-      annotations.push({
-        type: MessageAnnotationType.IMAGE,
-        data: { url: imageUrl },
-      });
-    }
-    if (files.length > 0) {
-      annotations.push({
-        type: MessageAnnotationType.DOCUMENT_FILE,
-        data: { files },
-      });
-    }
-    return annotations as JSONValue[];
-  };
-
-  const readContent = async (input: {
-    file: File;
-    asUrl?: boolean;
-  }): Promise<string> => {
-    const { file, asUrl } = input;
-    const content = await new Promise<string>((resolve, reject) => {
-      const reader = new FileReader();
-      if (asUrl) {
-        reader.readAsDataURL(file);
-      } else {
-        reader.readAsText(file);
-      }
-      reader.onload = () => resolve(reader.result as string);
-      reader.onerror = (error) => reject(error);
-    });
-    return content;
-  };
-
-  const uploadFile = async (file: File, requestParams: any = {}) => {
-    if (file.type.startsWith("image/")) {
-      const base64 = await readContent({ file, asUrl: true });
-      return setImageUrl(base64);
-    }
-
-    const filetype = docMineTypeMap[file.type];
-    if (!filetype) throw new Error("Unsupported document type.");
-    const newDoc = await uploadContent(file, requestParams);
-    return addDoc(newDoc);
-  };
-
-  return {
-    imageUrl,
-    setImageUrl,
-    files,
-    removeDoc,
-    reset,
-    getAnnotations,
-    uploadFile,
-  };
-}

templates/types/streaming/nextjs/app/components/ui/chat/index.ts

@@ -1,139 +0,0 @@
-import { JSONValue } from "ai";
-import ChatInput from "./chat-input";
-import ChatMessages from "./chat-messages";
-
-export { type ChatHandler } from "./chat.interface";
-export { ChatInput, ChatMessages };
-
-export enum MessageAnnotationType {
-  IMAGE = "image",
-  DOCUMENT_FILE = "document_file",
-  SOURCES = "sources",
-  EVENTS = "events",
-  TOOLS = "tools",
-  SUGGESTED_QUESTIONS = "suggested_questions",
-  AGENT_EVENTS = "agent",
-}
-
-export type ImageData = {
-  url: string;
-};
-
-export type DocumentFileType = "csv" | "pdf" | "txt" | "docx";
-export const DOCUMENT_FILE_TYPES: DocumentFileType[] = [
-  "csv",
-  "pdf",
-  "txt",
-  "docx",
-];
-
-export type DocumentFile = {
-  id: string;
-  name: string; // The uploaded file name in the backend
-  size: number; // The file size in bytes
-  type: DocumentFileType;
-  url: string; // The URL of the uploaded file in the backend
-  refs?: string[]; // DocumentIDs of the uploaded file in the vector index
-};
-
-export type DocumentFileData = {
-  files: DocumentFile[];
-};
-
-export type SourceNode = {
-  id: string;
-  metadata: Record<string, unknown>;
-  score?: number;
-  text: string;
-  url: string;
-};
-
-export type SourceData = {
-  nodes: SourceNode[];
-};
-
-export type EventData = {
-  title: string;
-};
-
-export type ProgressData = {
-  id: string;
-  total: number;
-  current: number;
-};
-
-export type AgentEventData = {
-  agent: string;
-  text: string;
-  type: "text" | "progress";
-  data?: ProgressData;
-};
-
-export type ToolData = {
-  toolCall: {
-    id: string;
-    name: string;
-    input: {
-      [key: string]: JSONValue;
-    };
-  };
-  toolOutput: {
-    output: JSONValue;
-    isError: boolean;
-  };
-};
-
-export type SuggestedQuestionsData = string[];
-
-export type AnnotationData =
-  | ImageData
-  | DocumentFileData
-  | SourceData
-  | EventData
-  | AgentEventData
-  | ToolData
-  | SuggestedQuestionsData;
-
-export type MessageAnnotation = {
-  type: MessageAnnotationType;
-  data: AnnotationData;
-};
-
-const NODE_SCORE_THRESHOLD = 0.25;
-
-export function getAnnotationData<T extends AnnotationData>(
-  annotations: MessageAnnotation[],
-  type: MessageAnnotationType,
-): T[] {
-  return annotations.filter((a) => a.type === type).map((a) => a.data as T);
-}
-
-export function getSourceAnnotationData(
-  annotations: MessageAnnotation[],
-): SourceData[] {
-  const data = getAnnotationData<SourceData>(
-    annotations,
-    MessageAnnotationType.SOURCES,
-  );
-  if (data.length > 0) {
-    const sourceData = data[0] as SourceData;
-    if (sourceData.nodes) {
-      sourceData.nodes = preprocessSourceNodes(sourceData.nodes);
-    }
-  }
-  return data;
-}
-
-function preprocessSourceNodes(nodes: SourceNode[]): SourceNode[] {
-  // Filter source nodes has lower score
-  nodes = nodes
-    .filter((node) => (node.score ?? 1) > NODE_SCORE_THRESHOLD)
-    .filter((node) => node.url && node.url.trim() !== "")
-    .sort((a, b) => (b.score ?? 1) - (a.score ?? 1))
-    .map((node) => {
-      // remove trailing slash for node url if exists
-      node.url = node.url.replace(/\/$/, "");
-      return node;
-    });
-  return nodes;
-}

templates/types/streaming/nextjs/app/components/ui/chat/tools/artifact.tsx

@@ -10,7 +10,7 @@ import {
 } from "../../collapsible";
 import { cn } from "../../lib/utils";
 import { Tabs, TabsContent, TabsList, TabsTrigger } from "../../tabs";
-import Markdown from "../chat-message/markdown";
+import { Markdown } from "../custom/markdown";
 import { useClientConfig } from "../hooks/use-config";
 import { useCopyToClipboard } from "../hooks/use-copy-to-clipboard";
 
@@ -29,12 +29,17 @@ export type CodeArtifact = {
   files?: string[];
 };
 
+type OutputUrl = {
+  url: string;
+  filename: string;
+};
+
 type ArtifactResult = {
   template: string;
   stdout: string[];
   stderr: string[];
   runtimeError?: { name: string; value: string; tracebackRaw: string[] };
-  outputUrls: Array<{ url: string; filename: string }>;
+  outputUrls: OutputUrl[];
   url: string;
 };
 
@@ -272,11 +277,7 @@ function CodeSandboxPreview({ url }: { url: string }) {
   );
 }
 
-function InterpreterOutput({
-  outputUrls,
-}: {
-  outputUrls: Array<{ url: string; filename: string }>;
-}) {
+function InterpreterOutput({ outputUrls }: { outputUrls: OutputUrl[] }) {
   return (
     <ul className="flex flex-col gap-2 mt-4">
       {outputUrls.map((url) => (

templates/types/streaming/nextjs/app/components/ui/chat/tools/chat-tools.tsx

@@ -0,0 +1,89 @@
+import {
+  getAnnotationData,
+  MessageAnnotation,
+  useChatMessage,
+  useChatUI,
+} from "@llamaindex/chat-ui";
+import { JSONValue, Message } from "ai";
+import { useMemo } from "react";
+import { Artifact, CodeArtifact } from "./artifact";
+import { WeatherCard, WeatherData } from "./weather-card";
+
+export function ToolAnnotations({ message }: { message: Message }) {
+  const annotations = message.annotations as MessageAnnotation[] | undefined;
+  const toolData = annotations
+    ? (getAnnotationData(annotations, "tools") as unknown as ToolData[])
+    : null;
+  return toolData?.[0] ? <ChatTools data={toolData[0]} /> : null;
+}
+
+// TODO: Used to render outputs of tools. If needed, add more renderers here.
+function ChatTools({ data }: { data: ToolData }) {
+  const { messages } = useChatUI();
+  const { message } = useChatMessage();
+
+  // build a map of message id to artifact version
+  const artifactVersionMap = useMemo(() => {
+    const map = new Map<string, number | undefined>();
+    let versionIndex = 1;
+    messages.forEach((m) => {
+      m.annotations?.forEach((annotation: any) => {
+        if (
+          typeof annotation === "object" &&
+          annotation != null &&
+          "type" in annotation &&
+          annotation.type === "tools"
+        ) {
+          const data = annotation.data as ToolData;
+          if (data?.toolCall?.name === "artifact") {
+            map.set(m.id, versionIndex);
+            versionIndex++;
+          }
+        }
+      });
+    });
+    return map;
+  }, [messages]);
+
+  if (!data) return null;
+  const { toolCall, toolOutput } = data;
+
+  if (toolOutput.isError) {
+    return (
+      <div className="border-l-2 border-red-400 pl-2">
+        There was an error when calling the tool {toolCall.name} with input:{" "}
+        <br />
+        {JSON.stringify(toolCall.input)}
+      </div>
+    );
+  }
+
+  switch (toolCall.name) {
+    case "get_weather_information":
+      const weatherData = toolOutput.output as unknown as WeatherData;
+      return <WeatherCard data={weatherData} />;
+    case "artifact":
+      return (
+        <Artifact
+          artifact={toolOutput.output as CodeArtifact}
+          version={artifactVersionMap.get(message.id)}
+        />
+      );
+    default:
+      return null;
+  }
+}
+
+type ToolData = {
+  toolCall: {
+    id: string;
+    name: string;
+    input: {
+      [key: string]: JSONValue;
+    };
+  };
+  toolOutput: {
+    output: JSONValue;
+    isError: boolean;
+  };
+};

templates/types/streaming/nextjs/app/components/ui/chat/widgets/PdfDialog.tsx

@@ -1,67 +0,0 @@
-import dynamic from "next/dynamic";
-import { Button } from "../../button";
-import {
-  Drawer,
-  DrawerClose,
-  DrawerContent,
-  DrawerDescription,
-  DrawerHeader,
-  DrawerTitle,
-  DrawerTrigger,
-} from "../../drawer";
-
-export interface PdfDialogProps {
-  documentId: string;
-  url: string;
-  trigger: React.ReactNode;
-}
-
-// Dynamic imports for client-side rendering only
-const PDFViewer = dynamic(
-  () => import("@llamaindex/pdf-viewer").then((module) => module.PDFViewer),
-  { ssr: false },
-);
-
-const PdfFocusProvider = dynamic(
-  () =>
-    import("@llamaindex/pdf-viewer").then((module) => module.PdfFocusProvider),
-  { ssr: false },
-);
-
-export default function PdfDialog(props: PdfDialogProps) {
-  return (
-    <Drawer direction="left">
-      <DrawerTrigger asChild>{props.trigger}</DrawerTrigger>
-      <DrawerContent className="w-3/5 mt-24 h-full max-h-[96%] ">
-        <DrawerHeader className="flex justify-between">
-          <div className="space-y-2">
-            <DrawerTitle>PDF Content</DrawerTitle>
-            <DrawerDescription>
-              File URL:{" "}
-              <a
-                className="hover:text-blue-900"
-                href={props.url}
-                target="_blank"
-              >
-                {props.url}
-              </a>
-            </DrawerDescription>
-          </div>
-          <DrawerClose asChild>
-            <Button variant="outline">Close</Button>
-          </DrawerClose>
-        </DrawerHeader>
-        <div className="m-4">
-          <PdfFocusProvider>
-            <PDFViewer
-              file={{
-                id: props.documentId,
-                url: props.url,
-              }}
-            />
-          </PdfFocusProvider>
-        </div>
-      </DrawerContent>
-    </Drawer>
-  );
-}

templates/types/streaming/nextjs/app/components/ui/document-preview.tsx

@@ -1,129 +0,0 @@
-import { XCircleIcon } from "lucide-react";
-import Image from "next/image";
-import DocxIcon from "../ui/icons/docx.svg";
-import PdfIcon from "../ui/icons/pdf.svg";
-import SheetIcon from "../ui/icons/sheet.svg";
-import TxtIcon from "../ui/icons/txt.svg";
-import { Button } from "./button";
-import { DocumentFile, DocumentFileType } from "./chat";
-import {
-  Drawer,
-  DrawerClose,
-  DrawerContent,
-  DrawerDescription,
-  DrawerHeader,
-  DrawerTitle,
-  DrawerTrigger,
-} from "./drawer";
-import { cn } from "./lib/utils";
-
-export interface DocumentPreviewProps {
-  file: DocumentFile;
-  onRemove?: () => void;
-}
-
-export function DocumentPreview(props: DocumentPreviewProps) {
-  const { name, size, type, refs } = props.file;
-
-  if (refs?.length) {
-    return (
-      <div title={`Document IDs: ${refs.join(", ")}`}>
-        <PreviewCard {...props} />
-      </div>
-    );
-  }
-
-  return (
-    <Drawer direction="left">
-      <DrawerTrigger asChild>
-        <div>
-          <PreviewCard className="cursor-pointer" {...props} />
-        </div>
-      </DrawerTrigger>
-      <DrawerContent className="w-3/5 mt-24 h-full max-h-[96%] ">
-        <DrawerHeader className="flex justify-between">
-          <div className="space-y-2">
-            <DrawerTitle>{type.toUpperCase()} Raw Content</DrawerTitle>
-            <DrawerDescription>
-              {name} ({inKB(size)} KB)
-            </DrawerDescription>
-          </div>
-          <DrawerClose asChild>
-            <Button variant="outline">Close</Button>
-          </DrawerClose>
-        </DrawerHeader>
-        <div className="m-4 max-h-[80%] overflow-auto">
-          {refs?.length && (
-            <pre className="bg-secondary rounded-md p-4 block text-sm">
-              {refs.join(", ")}
-            </pre>
-          )}
-        </div>
-      </DrawerContent>
-    </Drawer>
-  );
-}
-
-export const FileIcon: Record<DocumentFileType, string> = {
-  csv: SheetIcon,
-  pdf: PdfIcon,
-  docx: DocxIcon,
-  txt: TxtIcon,
-};
-
-export function PreviewCard(props: {
-  file: {
-    name: string;
-    size?: number;
-    type: DocumentFileType;
-  };
-  onRemove?: () => void;
-  className?: string;
-}) {
-  const { onRemove, file, className } = props;
-  return (
-    <div
-      className={cn(
-        "p-2 w-60 max-w-60 bg-secondary rounded-lg text-sm relative",
-        className,
-      )}
-    >
-      <div className="flex flex-row items-center gap-2">
-        <div className="relative h-8 w-8 shrink-0 overflow-hidden rounded-md flex items-center justify-center">
-          <Image
-            className="h-full w-auto object-contain"
-            priority
-            src={FileIcon[file.type]}
-            alt="Icon"
-          />
-        </div>
-        <div className="overflow-hidden">
-          <div className="truncate font-semibold">
-            {file.name} {file.size ? `(${inKB(file.size)} KB)` : ""}
-          </div>
-          {file.type && (
-            <div className="truncate text-token-text-tertiary flex items-center gap-2">
-              <span>{file.type.toUpperCase()} File</span>
-            </div>
-          )}
-        </div>
-      </div>
-      {onRemove && (
-        <div
-          className={cn(
-            "absolute -top-2 -right-2 w-6 h-6 z-10 bg-gray-500 text-white rounded-full",
-          )}
-        >
-          <XCircleIcon
-            className="w-6 h-6 bg-gray-500 text-white rounded-full"
-            onClick={onRemove}
-          />
-        </div>
-      )}
-    </div>
-  );
-}
-
-function inKB(size: number) {
-  return Math.round((size / 1024) * 10) / 10;
-}

templates/types/streaming/nextjs/app/components/ui/file-uploader.tsx

@@ -1,136 +0,0 @@
-"use client";
-
-import { Loader2, Paperclip } from "lucide-react";
-import { ChangeEvent, useState } from "react";
-import { buttonVariants } from "./button";
-import { cn } from "./lib/utils";
-
-export interface FileUploaderProps {
-  config?: {
-    inputId?: string;
-    fileSizeLimit?: number;
-    allowedExtensions?: string[];
-    checkExtension?: (extension: string) => string | null;
-    disabled: boolean;
-    multiple?: boolean;
-  };
-  onFileUpload: (file: File) => Promise<void>;
-  onFileError?: (errMsg: string) => void;
-}
-
-const DEFAULT_INPUT_ID = "fileInput";
-const DEFAULT_FILE_SIZE_LIMIT = 1024 * 1024 * 50; // 50 MB
-
-export default function FileUploader({
-  config,
-  onFileUpload,
-  onFileError,
-}: FileUploaderProps) {
-  const [uploading, setUploading] = useState(false);
-  const [remainingFiles, setRemainingFiles] = useState<number>(0);
-
-  const inputId = config?.inputId || DEFAULT_INPUT_ID;
-  const fileSizeLimit = config?.fileSizeLimit || DEFAULT_FILE_SIZE_LIMIT;
-  const allowedExtensions = config?.allowedExtensions;
-  const defaultCheckExtension = (extension: string) => {
-    if (allowedExtensions && !allowedExtensions.includes(extension)) {
-      return `Invalid file type. Please select a file with one of these formats: ${allowedExtensions!.join(
-        ",",
-      )}`;
-    }
-    return null;
-  };
-  const checkExtension = config?.checkExtension ?? defaultCheckExtension;
-
-  const isFileSizeExceeded = (file: File) => {
-    return file.size > fileSizeLimit;
-  };
-
-  const resetInput = () => {
-    const fileInput = document.getElementById(inputId) as HTMLInputElement;
-    fileInput.value = "";
-  };
-
-  const onFileChange = async (e: ChangeEvent<HTMLInputElement>) => {
-    const files = Array.from(e.target.files || []);
-    if (!files.length) return;
-
-    setUploading(true);
-
-    await handleUpload(files);
-
-    resetInput();
-    setUploading(false);
-  };
-
-  const handleUpload = async (files: File[]) => {
-    const onFileUploadError = onFileError || window.alert;
-    // Validate files
-    // If multiple files with image or multiple images
-    if (
-      files.length > 1 &&
-      files.some((file) => file.type.startsWith("image/"))
-    ) {
-      onFileUploadError("Multiple files with image are not supported");
-      return;
-    }
-
-    for (const file of files) {
-      const fileExtension = file.name.split(".").pop() || "";
-      const extensionFileError = checkExtension(fileExtension);
-      if (extensionFileError) {
-        onFileUploadError(extensionFileError);
-        return;
-      }
-
-      if (isFileSizeExceeded(file)) {
-        onFileUploadError(
-          `File size exceeded. Limit is ${fileSizeLimit / 1024 / 1024} MB`,
-        );
-        return;
-      }
-    }
-
-    setRemainingFiles(files.length);
-    for (const file of files) {
-      await onFileUpload(file);
-      setRemainingFiles((prev) => prev - 1);
-    }
-    setRemainingFiles(0);
-  };
-
-  return (
-    <div className="self-stretch">
-      <input
-        type="file"
-        id={inputId}
-        style={{ display: "none" }}
-        onChange={onFileChange}
-        accept={allowedExtensions?.join(",")}
-        disabled={config?.disabled || uploading}
-        multiple={config?.multiple}
-      />
-      <label
-        htmlFor={inputId}
-        className={cn(
-          buttonVariants({ variant: "secondary", size: "icon" }),
-          "cursor-pointer relative",
-          uploading && "opacity-50",
-        )}
-      >
-        {uploading ? (
-          <div className="relative flex items-center justify-center h-full w-full">
-            <Loader2 className="h-6 w-6 animate-spin absolute" />
-            {remainingFiles > 0 && (
-              <span className="text-xs absolute inset-0 flex items-center justify-center">
-                {remainingFiles}
-              </span>
-            )}
-          </div>
-        ) : (
-          <Paperclip className="-rotate-45 w-4 h-4" />
-        )}
-      </label>
-    </div>
-  );
-}

templates/types/streaming/nextjs/app/components/ui/upload-image-preview.tsx

@@ -1,32 +0,0 @@
-import { XCircleIcon } from "lucide-react";
-import Image from "next/image";
-import { cn } from "./lib/utils";
-
-export default function UploadImagePreview({
-  url,
-  onRemove,
-}: {
-  url: string;
-  onRemove: () => void;
-}) {
-  return (
-    <div className="relative w-20 h-20 group">
-      <Image
-        src={url}
-        alt="Uploaded image"
-        fill
-        className="object-cover w-full h-full rounded-xl hover:brightness-75"
-      />
-      <div
-        className={cn(
-          "absolute -top-2 -right-2 w-6 h-6 z-10 bg-gray-500 text-white rounded-full hidden group-hover:block",
-        )}
-      >
-        <XCircleIcon
-          className="w-6 h-6 bg-gray-500 text-white rounded-full"
-          onClick={onRemove}
-        />
-      </div>
-    </div>
-  );
-}

templates/types/streaming/nextjs/app/observability/index.ts

@@ -1 +1,2 @@
+// TODO: You can add observability here. For templates re-start `create-llama` with `--pro` flag to generate a new project with observability.
 export const initObservability = () => {};

templates/types/streaming/nextjs/package.json

@@ -12,7 +12,6 @@
   "dependencies": {
     "@apidevtools/swagger-parser": "^10.1.0",
     "@e2b/code-interpreter": "0.0.9-beta.3",
-    "@llamaindex/pdf-viewer": "^1.1.3",
     "@radix-ui/react-collapsible": "^1.0.3",
     "@radix-ui/react-hover-card": "^1.0.7",
     "@radix-ui/react-progress": "^1.1.0",
@@ -27,24 +26,18 @@
     "duck-duck-scrape": "^2.2.5",
     "formdata-node": "^6.0.3",
     "got": "^14.4.1",
-    "llamaindex": "0.7.10",
+    "llamaindex": "0.8.2",
     "lucide-react": "^0.294.0",
     "next": "^14.2.4",
     "react": "^18.2.0",
     "react-dom": "^18.2.0",
-    "react-markdown": "^8.0.7",
-    "rehype-katex": "^7.0.0",
-    "remark": "^14.0.3",
-    "remark-code-import": "^1.2.0",
-    "remark-gfm": "^3.0.1",
-    "remark-math": "^5.1.1",
     "supports-color": "^8.1.1",
     "tailwind-merge": "^2.1.0",
     "tiktoken": "^1.0.15",
     "uuid": "^9.0.1",
     "vaul": "^0.9.1",
     "marked": "^14.1.2",
-    "highlight.js": "^11.10.0"
+    "@llamaindex/chat-ui": "0.0.4"
   },
   "devDependencies": {
     "@types/node": "^20.10.3",

templates/types/streaming/nextjs/tailwind.config.ts

@@ -3,7 +3,11 @@ import { fontFamily } from "tailwindcss/defaultTheme";
 
 const config: Config = {
   darkMode: ["class"],
-  content: ["app/**/*.{ts,tsx}", "components/**/*.{ts,tsx}"],
+  content: [
+    "app/**/*.{ts,tsx}",
+    "components/**/*.{ts,tsx}",
+    "node_modules/@llamaindex/chat-ui/**/*.{ts,tsx}",
+  ],
   theme: {
     container: {
       center: true,