fix(useWorkflow): handle incomplete chunk and support tool call transformation (#159)

* fix: handle fail to parse chunk * add test * Create gorgeous-lies-develop.md * support transforming toolcall and toolcalloutput * raw node to source node * fileServerUrl option for useChatWorkflow * add normal agentic rag example * missing file * return slug url * update pnpm workspace * fix format * feat: add agentic rag cloud example * fix format * remove dead files * fix data * RawNodeWithScore * update url * add comment * log info * update llamacloud example
2026-07-22 03:45:20 -04:00 · 2025-07-09 14:41:50 +07:00
parent 4855f5ce71
commit 62955a6a18
37 changed files with 1580 additions and 37 deletions
@@ -0,0 +1,5 @@
+---
+'@llamaindex/chat-ui': patch
+---
+
+fix(useWorkflow): handle incomplete chunk from llama-deploy workflow and retry parsing
@@ -0,0 +1,5 @@
+__pycache__
+storage
+.env
+output
+.ui/
@@ -0,0 +1,106 @@
+# LlamaIndex Workflow Example
+
+This is a [LlamaIndex](https://www.llamaindex.ai/) project that using [Workflows](https://docs.llamaindex.ai/en/stable/understanding/workflows/) deployed with [LlamaDeploy](https://github.com/run-llama/llama_deploy).
+
+LlamaDeploy is a system for deploying and managing LlamaIndex workflows, while LlamaIndexServer provides a pre-built TypeScript server with an integrated chat UI that can connect directly to LlamaDeploy deployments. This example shows how you can quickly set up a complete chat application by combining these two technologies/
+
+## Prerequisites
+
+If you haven't installed uv, you can follow the instructions [here](https://docs.astral.sh/uv/getting-started/installation/) to install it.
+
+You can configure [LLM model](https://docs.llamaindex.ai/en/stable/module_guides/models/llms) and [embedding model](https://docs.llamaindex.ai/en/stable/module_guides/models/embeddings) in [src/settings.py](src/settings.py).
+
+Please setup their API keys in the `src/.env` file.
+
+## Installation
+
+Both the SDK and the CLI are part of the LlamaDeploy Python package. To install, just run:
+
+```bash
+uv sync
+```
+
+If you don't have uv installed, you can follow the instructions [here](https://docs.astral.sh/uv/getting-started/installation/).
+
+## Generate Index
+
+Generate the embeddings of the documents in the `./data` directory:
+
+```shell
+uv run generate
+```
+
+## Running the Deployment
+
+At this point we have all we need to run this deployment. Ideally, we would have the API server already running
+somewhere in the cloud, but to get started let's start an instance locally. Run the following python script
+from a shell:
+
+```
+$ uv run -m llama_deploy.apiserver
+INFO:     Started server process [10842]
+INFO:     Waiting for application startup.
+INFO:     Application startup complete.
+INFO:     Uvicorn running on http://0.0.0.0:4501 (Press CTRL+C to quit)
+```
+
+From another shell, use the CLI, `llamactl`, to create the deployment:
+
+```
+$ uv run llamactl deploy llama_deploy.yml
+Deployment successful: chat
+```
+
+## UI Interface
+
+LlamaDeploy will serve the UI through the apiserver. Point the browser to [http://localhost:4501/deployments/chat/ui](http://localhost:4501/deployments/chat/ui) to interact with your deployment through a user-friendly interface.
+
+## API endpoints
+
+You can find all the endpoints in the [API documentation](http://localhost:4501/docs). To get started, you can try the following endpoints:
+
+Create a new task:
+
+```bash
+curl -X POST 'http://localhost:4501/deployments/chat/tasks/create' \
+  -H 'Content-Type: application/json' \
+  -d '{
+    "input": "{\"user_msg\":\"Hello\",\"chat_history\":[]}",
+    "service_id": "workflow"
+  }'
+```
+
+Stream events:
+
+```bash
+curl 'http://localhost:4501/deployments/chat/tasks/0b411be6-005d-43f0-9b6b-6a0017f08002/events?session_id=dd36442c-45ca-4eaa-8d75-b4e6dad1a83e&raw_event=true' \
+  -H 'Content-Type: application/json'
+```
+
+Note that the task_id and session_id are returned when creating a new task.
+
+## Use Case
+
+We have prepared an [example workflow](./src/workflow.py) for the agentic RAG use case, where you can ask questions about the example documents in the [./data](./data) directory.
+To update the workflow, you can modify the code in [`src/workflow.py`](src/workflow.py).
+
+## Customize the UI
+
+The UI is served by LLamaIndexServer package, you can configure the UI by modifying the `uiConfig` in the [ui/index.ts](ui/index.ts) file.
+
+The following are the available options:
+
+- `starterQuestions`: Predefined questions for chat interface
+- `componentsDir`: Directory for custom event components
+- `layoutDir`: Directory for custom layout components
+- `llamaCloudIndexSelector`: Enable LlamaCloud integration
+- `llamaDeploy`: The LlamaDeploy configration (deployment name and workflow name that defined in the [llama_deploy.yml](llama_deploy.yml) file)
+
+## Learn More
+
+- [LlamaIndex Documentation](https://docs.llamaindex.ai) - learn about LlamaIndex.
+- [Workflows Introduction](https://docs.llamaindex.ai/en/stable/understanding/workflows/) - learn about LlamaIndex workflows.
+- [LlamaDeploy GitHub Repository](https://github.com/run-llama/llama_deploy)
+- [Chat-UI Documentation](https://ts.llamaindex.ai/docs/chat-ui)
+
+You can check out [the LlamaIndex GitHub repository](https://github.com/run-llama/llama_index) - your feedback and contributions are welcome!
@@ -0,0 +1,24 @@
+name: chat
+
+control-plane:
+  port: 8000
+
+default-service: workflow
+
+services:
+  workflow:
+    name: Workflow
+    source:
+      type: local
+      name: src
+    path: src/workflow:workflow
+    python-dependencies:
+      - llama-index-llms-openai>=0.4.5
+      - llama-index-core>=0.12.45
+
+ui:
+  name: My Nextjs App
+  port: 3000
+  source:
+    type: local
+    name: ui
@@ -0,0 +1,54 @@
+[project]
+name = "app"
+version = "0.1.0"
+description = ""
+readme = "README.md"
+requires-python = ">=3.11,<3.14"
+dependencies = [ "python-dotenv>=1.0.0,<2.0.0", "pydantic>=2.11.5", "aiostream>=0.5.2,<0.6.0", "llama-index-core>=0.12.28,<0.13.0", "llama-index-readers-file>=0.4.6,<1.0.0", "llama-index-indices-managed-llama-cloud>=0.6.3,<1.0.0", "llama-deploy", "llama-index-indices-managed-llama-cloud>=0.6.3,<0.7.0", "docx2txt>=0.8,<0.9", "llama-index-llms-openai>=0.3.2,<0.4.0", "llama-index-embeddings-openai>=0.3.1,<0.4.0" ]
+
+[[project.authors]]
+name = "Marcus Schiesser"
+email = "mail@marcusschiesser.de"
+
+[project.optional-dependencies]
+dev = [ "mypy>=1.8.0,<2.0.0", "pytest>=8.3.5,<9.0.0", "pytest-asyncio>=0.25.3,<0.26.0" ]
+
+[project.scripts]
+generate = "src.generate:generate_index"
+
+[tool]
+[tool.uv]
+[tool.uv.sources]
+[tool.uv.sources.llama-deploy]
+git = "https://github.com/run-llama/llama_deploy"
+
+[tool.mypy]
+python_version = "3.11"
+plugins = "pydantic.mypy"
+exclude = [ "tests", "venv", ".venv", "output", "config" ]
+check_untyped_defs = true
+warn_unused_ignores = false
+show_error_codes = true
+namespace_packages = true
+ignore_missing_imports = true
+follow_imports = "silent"
+implicit_optional = true
+strict_optional = false
+disable_error_code = [ "return-value", "assignment" ]
+
+[[tool.mypy.overrides]]
+module = "src.*"
+ignore_missing_imports = false
+
+[tool.hatch]
+[tool.hatch.metadata]
+allow-direct-references = true
+
+[tool.hatch.build]
+[tool.hatch.build.targets]
+[tool.hatch.build.targets.wheel]
+packages = [ "src" ]
+
+[build-system]
+requires = [ "hatchling>=1.24" ]
+build-backend = "hatchling.build"
@@ -0,0 +1,106 @@
+from typing import Any, List, Optional
+
+from llama_index.core import QueryBundle
+from llama_index.core.postprocessor.types import BaseNodePostprocessor
+from llama_index.core.prompts import PromptTemplate
+from llama_index.core.query_engine.retriever_query_engine import RetrieverQueryEngine
+from llama_index.core.response_synthesizers import Accumulate
+from llama_index.core.schema import NodeWithScore
+from llama_index.core.tools.query_engine import QueryEngineTool
+
+
+# Used as a prompt for synthesizer
+# Override this prompt by setting the `CITATION_PROMPT` environment variable
+CITATION_PROMPT = """
+Context information is below.
+------------------
+{context_str}
+------------------
+The context are multiple text chunks, each text chunk has its own citation_id at the beginning.
+Use the citation_id for citation construction.
+
+Answer the following query with citations:
+------------------
+{query_str}
+------------------
+
+## Citation format
+
+[citation:id]
+
+Where:
+- [citation:] is a matching pattern which is required for all citations.
+- `id` is the `citation_id` provided in the context or previous response.
+
+Example:
+```
+    Here is a response that uses context information [citation:90ca859f-4f32-40ca-8cd0-edfad4fb298b] 
+    and other ideas that don't use context information [citation:17b2cc9a-27ae-4b6d-bede-5ca60fc00ff4] .\n
+    The citation block will be displayed automatically with useful information for the user in the UI [citation:1c606612-e75f-490e-8374-44e79f818d19] .
+```
+
+## Requirements:
+1. Always include citations for every fact from the context information in your response. 
+2. Make sure that the citation_id is correct with the context, don't mix up the citation_id with other information.
+
+Now, you answer the query with citations:
+"""
+
+
+class NodeCitationProcessor(BaseNodePostprocessor):
+    """
+    Add a new field `citation_id` to the metadata of the node by copying the id from the node.
+    Useful for citation construction.
+    """
+
+    def _postprocess_nodes(
+        self,
+        nodes: List[NodeWithScore],
+        query_bundle: Optional[QueryBundle] = None,
+    ) -> List[NodeWithScore]:
+        for node_score in nodes:
+            node_score.node.metadata["citation_id"] = node_score.node.node_id
+        return nodes
+
+
+class CitationSynthesizer(Accumulate):
+    """
+    Overload the Accumulate synthesizer to:
+    1. Update prepare node metadata for citation id
+    2. Update text_qa_template to include citations
+    """
+
+    def __init__(self, **kwargs: Any) -> None:
+        text_qa_template = kwargs.pop("text_qa_template", None)
+        if text_qa_template is None:
+            text_qa_template = PromptTemplate(template=CITATION_PROMPT)
+        super().__init__(text_qa_template=text_qa_template, **kwargs)
+
+
+# Add this prompt to your agent system prompt
+CITATION_SYSTEM_PROMPT = (
+    "\nAnswer the user question using the response from the query tool. "
+    "It's important to respect the citation information in the response. "
+    "Don't mix up the citation_id, keep them at the correct fact."
+)
+
+
+def enable_citation(query_engine_tool: QueryEngineTool) -> QueryEngineTool:
+    """
+    Enable citation for a query engine tool by using CitationSynthesizer and NodePostprocessor.
+    Note: This function will override the response synthesizer of your query engine.
+    """
+    query_engine = query_engine_tool.query_engine
+    if not isinstance(query_engine, RetrieverQueryEngine):
+        raise ValueError(
+            "Citation feature requires a RetrieverQueryEngine. Your tool's query engine is a "
+            f"{type(query_engine)}."
+        )
+    # Update the response synthesizer and node postprocessors
+    query_engine._response_synthesizer = CitationSynthesizer()
+    query_engine._node_postprocessors += [NodeCitationProcessor()]
+    query_engine_tool._query_engine = query_engine
+
+    # Update tool metadata
+    query_engine_tool.metadata.description += "\nThe output will include citations with the format [citation:id] for each chunk of information in the knowledge base."
+    return query_engine_tool
@@ -0,0 +1,65 @@
+# flake8: noqa: E402
+
+from dotenv import load_dotenv
+
+load_dotenv()
+
+import logging
+
+from llama_index.core.readers import SimpleDirectoryReader
+from tqdm import tqdm
+
+from src.index import get_index
+from src.service import LLamaCloudFileService
+from src.settings import init_settings
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger()
+
+
+def generate_index():
+    init_settings()
+    logger.info("Generate index for the provided data")
+
+    index = get_index(create_if_missing=True)
+    if index is None:
+        raise ValueError("Index not found and could not be created")
+
+    # use SimpleDirectoryReader to retrieve the files to process
+    reader = SimpleDirectoryReader(
+        "ui/data",
+        recursive=True,
+    )
+    files_to_process = reader.input_files
+
+    # add each file to the LlamaCloud pipeline
+    error_files = []
+    for input_file in tqdm(
+        files_to_process,
+        desc="Processing files",
+        unit="file",
+    ):
+        with open(input_file, "rb") as f:
+            logger.debug(
+                f"Adding file {input_file} to pipeline {index.name} in project {index.project_name}"
+            )
+            try:
+                LLamaCloudFileService.add_file_to_pipeline(
+                    index.project.id,
+                    index.pipeline.id,
+                    f,
+                    custom_metadata={},
+                    wait_for_processing=False,
+                )
+            except Exception as e:
+                error_files.append(input_file)
+                logger.error(f"Error adding file {input_file}: {e}")
+
+    if error_files:
+        logger.error(f"Failed to add the following files: {error_files}")
+
+    logger.info("Finished generating the index")
+
+
+if __name__ == "__main__":
+    generate_index()
@@ -0,0 +1,146 @@
+import logging
+import os
+from typing import Optional
+
+from llama_cloud import PipelineType
+from llama_index.core.callbacks import CallbackManager
+from llama_index.core.ingestion.api_utils import (
+    get_client as llama_cloud_get_client,
+)
+from llama_index.core.settings import Settings
+from llama_index.indices.managed.llama_cloud import LlamaCloudIndex
+from pydantic import BaseModel, Field, field_validator
+
+logger = logging.getLogger("uvicorn")
+
+
+class LlamaCloudConfig(BaseModel):
+    # Private attributes
+    api_key: str = Field(
+        exclude=True,  # Exclude from the model representation
+    )
+    base_url: Optional[str] = Field(
+        exclude=True,
+    )
+    organization_id: Optional[str] = Field(
+        exclude=True,
+    )
+    # Configuration attributes, can be set by the user
+    pipeline: str = Field(
+        description="The name of the pipeline to use",
+    )
+    project: str = Field(
+        description="The name of the LlamaCloud project",
+    )
+
+    def __init__(self, **kwargs):
+        if "api_key" not in kwargs:
+            kwargs["api_key"] = os.getenv("LLAMA_CLOUD_API_KEY")
+        if "base_url" not in kwargs:
+            kwargs["base_url"] = os.getenv("LLAMA_CLOUD_BASE_URL")
+        if "organization_id" not in kwargs:
+            kwargs["organization_id"] = os.getenv("LLAMA_CLOUD_ORGANIZATION_ID")
+        if "pipeline" not in kwargs:
+            kwargs["pipeline"] = os.getenv("LLAMA_CLOUD_INDEX_NAME")
+        if "project" not in kwargs:
+            kwargs["project"] = os.getenv("LLAMA_CLOUD_PROJECT_NAME")
+        super().__init__(**kwargs)
+
+    # Validate and throw error if the env variables are not set before starting the app
+    @field_validator("pipeline", "project", "api_key", mode="before")
+    @classmethod
+    def validate_fields(cls, value):
+        if value is None:
+            raise ValueError(
+                "Please set LLAMA_CLOUD_INDEX_NAME, LLAMA_CLOUD_PROJECT_NAME and LLAMA_CLOUD_API_KEY"
+                " to your environment variables or config them in .env file"
+            )
+        return value
+
+    def to_client_kwargs(self) -> dict:
+        return {
+            "api_key": self.api_key,
+            "base_url": self.base_url,
+        }
+
+
+class IndexConfig(BaseModel):
+    llama_cloud_pipeline_config: LlamaCloudConfig = Field(
+        default_factory=lambda: LlamaCloudConfig(),
+        alias="llamaCloudPipeline",
+    )
+    callback_manager: Optional[CallbackManager] = Field(
+        default=None,
+    )
+
+    def to_index_kwargs(self) -> dict:
+        return {
+            "name": self.llama_cloud_pipeline_config.pipeline,
+            "project_name": self.llama_cloud_pipeline_config.project,
+            "api_key": self.llama_cloud_pipeline_config.api_key,
+            "base_url": self.llama_cloud_pipeline_config.base_url,
+            "organization_id": self.llama_cloud_pipeline_config.organization_id,
+            "callback_manager": self.callback_manager,
+        }
+
+
+def get_index(
+    config: IndexConfig = None,
+    create_if_missing: bool = False,
+):
+    if config is None:
+        config = IndexConfig()
+    # Check whether the index exists
+    try:
+        index = LlamaCloudIndex(**config.to_index_kwargs())
+        return index
+    except ValueError:
+        logger.warning("Index not found")
+        if create_if_missing:
+            logger.info("Creating index")
+            _create_index(config)
+            return LlamaCloudIndex(**config.to_index_kwargs())
+        return None
+
+
+def get_client():
+    config = LlamaCloudConfig()
+    return llama_cloud_get_client(**config.to_client_kwargs())
+
+
+def _create_index(
+    config: IndexConfig,
+):
+    client = get_client()
+    pipeline_name = config.llama_cloud_pipeline_config.pipeline
+
+    pipelines = client.pipelines.search_pipelines(
+        pipeline_name=pipeline_name,
+        pipeline_type=PipelineType.MANAGED.value,
+    )
+    if len(pipelines) == 0:
+        from llama_index.embeddings.openai import OpenAIEmbedding
+
+        if not isinstance(Settings.embed_model, OpenAIEmbedding):
+            raise ValueError(
+                "Creating a new pipeline with a non-OpenAI embedding model is not supported."
+            )
+        client.pipelines.upsert_pipeline(
+            request={
+                "name": pipeline_name,
+                "embedding_config": {
+                    "type": "OPENAI_EMBEDDING",
+                    "component": {
+                        "api_key": os.getenv("OPENAI_API_KEY"),  # editable
+                        "model_name": os.getenv("EMBEDDING_MODEL"),
+                    },
+                },
+                "transform_config": {
+                    "mode": "auto",
+                    "config": {
+                        "chunk_size": Settings.chunk_size,  # editable
+                        "chunk_overlap": Settings.chunk_overlap,  # editable
+                    },
+                },
+            },
+        )
@@ -0,0 +1,47 @@
+import os
+from typing import Any, Optional
+
+from llama_index.core.base.base_query_engine import BaseQueryEngine
+from llama_index.core.indices.base import BaseIndex
+from llama_index.core.tools.query_engine import QueryEngineTool
+
+def create_query_engine(index: BaseIndex, **kwargs: Any) -> BaseQueryEngine:
+    """
+    Create a query engine for the given index.
+
+    Args:
+        index: The index to create a query engine for.
+        params (optional): Additional parameters for the query engine, e.g: similarity_top_k
+    """
+    top_k = int(os.getenv("TOP_K", 0))
+    if top_k != 0 and kwargs.get("filters") is None:
+        kwargs["similarity_top_k"] = top_k
+
+    return index.as_query_engine(**kwargs)
+
+
+def get_query_engine_tool(
+    index: BaseIndex,
+    name: Optional[str] = None,
+    description: Optional[str] = None,
+    **kwargs: Any,
+) -> QueryEngineTool:
+    """
+    Get a query engine tool for the given index.
+
+    Args:
+        index: The index to create a query engine for.
+        name (optional): The name of the tool.
+        description (optional): The description of the tool.
+    """
+    if name is None:
+        name = "query_index"
+    if description is None:
+        description = "Use this tool to retrieve information from a knowledge base. Provide a specific query and can call the tool multiple times if necessary."
+    query_engine = create_query_engine(index, **kwargs)
+    tool = QueryEngineTool.from_defaults(
+        query_engine=query_engine,
+        name=name,
+        description=description,
+    )
+    return tool
@@ -0,0 +1,74 @@
+import logging
+import os
+import time
+import typing
+from io import BytesIO
+from typing import Any, Dict, List, Optional, Set, Tuple, Union
+
+from llama_cloud import ManagedIngestionStatus, PipelineFileCreateCustomMetadataValue
+from pydantic import BaseModel
+
+from src.index import get_client
+
+logger = logging.getLogger("uvicorn")
+
+
+class LlamaCloudFile(BaseModel):
+    file_name: str
+    pipeline_id: str
+
+    def __eq__(self, other):
+        if not isinstance(other, LlamaCloudFile):
+            return NotImplemented
+        return (
+            self.file_name == other.file_name and self.pipeline_id == other.pipeline_id
+        )
+
+    def __hash__(self):
+        return hash((self.file_name, self.pipeline_id))
+
+
+class LLamaCloudFileService:
+    LOCAL_STORE_PATH = "output/llamacloud"
+    DOWNLOAD_FILE_NAME_TPL = "{pipeline_id}${filename}"
+
+    @classmethod
+    def add_file_to_pipeline(
+        cls,
+        project_id: str,
+        pipeline_id: str,
+        upload_file: Union[typing.IO, Tuple[str, BytesIO]],
+        custom_metadata: Optional[Dict[str, PipelineFileCreateCustomMetadataValue]],
+        wait_for_processing: bool = True,
+    ) -> str:
+        client = get_client()
+        file = client.files.upload_file(project_id=project_id, upload_file=upload_file)
+        file_id = file.id
+        files = [
+            {
+                "file_id": file_id,
+                "custom_metadata": {"file_id": file_id, **(custom_metadata or {})},
+            }
+        ]
+        files = client.pipelines.add_files_to_pipeline_api(pipeline_id, request=files)
+
+        if not wait_for_processing:
+            return file_id
+
+        # Wait 2s for the file to be processed
+        max_attempts = 20
+        attempt = 0
+        while attempt < max_attempts:
+            result = client.pipelines.get_pipeline_file_status(
+                file_id=file_id, pipeline_id=pipeline_id
+            )
+            if result.status == ManagedIngestionStatus.ERROR:
+                raise Exception(f"File processing failed: {str(result)}")
+            if result.status == ManagedIngestionStatus.SUCCESS:
+                # File is ingested - return the file id
+                return file_id
+            attempt += 1
+            time.sleep(0.1)  # Sleep for 100ms
+        raise Exception(
+            f"File processing did not complete after {max_attempts} attempts."
+        )
@@ -0,0 +1,14 @@
+import os
+
+from llama_index.core import Settings
+from llama_index.embeddings.openai import OpenAIEmbedding
+from llama_index.llms.openai import OpenAI
+
+
+def init_settings():
+    if os.getenv("OPENAI_API_KEY") is None:
+        raise RuntimeError("OPENAI_API_KEY is missing in environment variables")
+    Settings.llm = OpenAI(model=os.getenv("MODEL") or "gpt-4.1")
+    Settings.embed_model = OpenAIEmbedding(
+        model=os.getenv("EMBEDDING_MODEL") or "text-embedding-3-large"
+    )
@@ -0,0 +1,35 @@
+from dotenv import load_dotenv
+
+from llama_index.core.agent.workflow import AgentWorkflow
+from llama_index.core.settings import Settings
+
+from src.index import get_index
+from src.query import get_query_engine_tool
+from src.citation import CITATION_SYSTEM_PROMPT, enable_citation
+from src.settings import init_settings
+
+
+def create_workflow() -> AgentWorkflow:
+    load_dotenv()
+    init_settings()
+    index = get_index()
+    if index is None:
+        raise RuntimeError(
+            "Index not found! Please run `uv run generate` to index the data first."
+        )
+    # Create a query tool with citations enabled
+    query_tool = enable_citation(get_query_engine_tool(index=index))
+
+    # Define the system prompt for the agent
+    # Append the citation system prompt to the system prompt
+    system_prompt = """You are a helpful assistant"""
+    system_prompt += CITATION_SYSTEM_PROMPT
+
+    return AgentWorkflow.from_tools_or_functions(
+        tools_or_functions=[query_tool],
+        llm=Settings.llm,
+        system_prompt=system_prompt,
+    )
+
+
+workflow = create_workflow()
@@ -0,0 +1,10 @@
+import { LlamaIndexServer } from '@llamaindex/server'
+
+new LlamaIndexServer({
+  uiConfig: {
+    componentsDir: 'components',
+    layoutDir: 'layout',
+    llamaDeploy: { deployment: 'chat', workflow: 'workflow' },
+  },
+  port: 3000,
+}).start()
@@ -0,0 +1,40 @@
+'use client'
+
+import { Sparkles, Star } from 'lucide-react'
+
+export default function Header() {
+  return (
+    <div className="flex items-center justify-between p-2 px-4">
+      <div className="flex items-center gap-2">
+        <Sparkles className="size-4" />
+        <h1 className="font-semibold">LlamaIndex App</h1>
+      </div>
+      <div className="flex items-center justify-end gap-4">
+        <div className="flex items-center gap-2">
+          <a
+            href="https://www.llamaindex.ai/"
+            target="_blank"
+            rel="noopener noreferrer"
+            className="text-sm text-gray-600 hover:text-gray-800 dark:text-gray-400 dark:hover:text-gray-200"
+          >
+            Built by LlamaIndex
+          </a>
+          <img
+            className="h-[24px] w-[24px] rounded-sm"
+            src="https://ui.llamaindex.ai/llama.png"
+            alt="Llama Logo"
+          />
+        </div>
+        <a
+          href="https://github.com/run-llama/LlamaIndexTS"
+          target="_blank"
+          rel="noopener noreferrer"
+          className="hover:bg-accent flex items-center gap-2 rounded-md border border-gray-300 px-2 py-1 text-sm"
+        >
+          <Star className="size-4" />
+          Star on GitHub
+        </a>
+      </div>
+    </div>
+  )
+}
@@ -0,0 +1,18 @@
+{
+  "name": "llamaindex-server-ui",
+  "version": "0.0.1",
+  "private": true,
+  "scripts": {
+    "dev": "nodemon --exec tsx index.ts"
+  },
+  "dependencies": {
+    "@llamaindex/server": "latest",
+    "dotenv": "^16.4.7"
+  },
+  "devDependencies": {
+    "@types/node": "^20.10.3",
+    "nodemon": "^3.1.10",
+    "tsx": "4.7.2",
+    "typescript": "^5.3.2"
+  }
+}
@@ -0,0 +1,5 @@
+__pycache__
+storage
+.env
+output
+.ui/
@@ -0,0 +1,106 @@
+# LlamaIndex Workflow Example
+
+This is a [LlamaIndex](https://www.llamaindex.ai/) project that using [Workflows](https://docs.llamaindex.ai/en/stable/understanding/workflows/) deployed with [LlamaDeploy](https://github.com/run-llama/llama_deploy).
+
+LlamaDeploy is a system for deploying and managing LlamaIndex workflows, while LlamaIndexServer provides a pre-built TypeScript server with an integrated chat UI that can connect directly to LlamaDeploy deployments. This example shows how you can quickly set up a complete chat application by combining these two technologies/
+
+## Prerequisites
+
+If you haven't installed uv, you can follow the instructions [here](https://docs.astral.sh/uv/getting-started/installation/) to install it.
+
+You can configure [LLM model](https://docs.llamaindex.ai/en/stable/module_guides/models/llms) and [embedding model](https://docs.llamaindex.ai/en/stable/module_guides/models/embeddings) in [src/settings.py](src/settings.py).
+
+Please setup their API keys in the `src/.env` file.
+
+## Installation
+
+Both the SDK and the CLI are part of the LlamaDeploy Python package. To install, just run:
+
+```bash
+uv sync
+```
+
+If you don't have uv installed, you can follow the instructions [here](https://docs.astral.sh/uv/getting-started/installation/).
+
+## Generate Index
+
+Generate the embeddings of the documents in the `./data` directory:
+
+```shell
+uv run generate
+```
+
+## Running the Deployment
+
+At this point we have all we need to run this deployment. Ideally, we would have the API server already running
+somewhere in the cloud, but to get started let's start an instance locally. Run the following python script
+from a shell:
+
+```
+$ uv run -m llama_deploy.apiserver
+INFO:     Started server process [10842]
+INFO:     Waiting for application startup.
+INFO:     Application startup complete.
+INFO:     Uvicorn running on http://0.0.0.0:4501 (Press CTRL+C to quit)
+```
+
+From another shell, use the CLI, `llamactl`, to create the deployment:
+
+```
+$ uv run llamactl deploy llama_deploy.yml
+Deployment successful: chat
+```
+
+## UI Interface
+
+LlamaDeploy will serve the UI through the apiserver. Point the browser to [http://localhost:4501/deployments/chat/ui](http://localhost:4501/deployments/chat/ui) to interact with your deployment through a user-friendly interface.
+
+## API endpoints
+
+You can find all the endpoints in the [API documentation](http://localhost:4501/docs). To get started, you can try the following endpoints:
+
+Create a new task:
+
+```bash
+curl -X POST 'http://localhost:4501/deployments/chat/tasks/create' \
+  -H 'Content-Type: application/json' \
+  -d '{
+    "input": "{\"user_msg\":\"Hello\",\"chat_history\":[]}",
+    "service_id": "workflow"
+  }'
+```
+
+Stream events:
+
+```bash
+curl 'http://localhost:4501/deployments/chat/tasks/0b411be6-005d-43f0-9b6b-6a0017f08002/events?session_id=dd36442c-45ca-4eaa-8d75-b4e6dad1a83e&raw_event=true' \
+  -H 'Content-Type: application/json'
+```
+
+Note that the task_id and session_id are returned when creating a new task.
+
+## Use Case
+
+We have prepared an [example workflow](./src/workflow.py) for the agentic RAG use case, where you can ask questions about the example documents in the [./data](./data) directory.
+To update the workflow, you can modify the code in [`src/workflow.py`](src/workflow.py).
+
+## Customize the UI
+
+The UI is served by LLamaIndexServer package, you can configure the UI by modifying the `uiConfig` in the [ui/index.ts](ui/index.ts) file.
+
+The following are the available options:
+
+- `starterQuestions`: Predefined questions for chat interface
+- `componentsDir`: Directory for custom event components
+- `layoutDir`: Directory for custom layout components
+- `llamaCloudIndexSelector`: Enable LlamaCloud integration
+- `llamaDeploy`: The LlamaDeploy configration (deployment name and workflow name that defined in the [llama_deploy.yml](llama_deploy.yml) file)
+
+## Learn More
+
+- [LlamaIndex Documentation](https://docs.llamaindex.ai) - learn about LlamaIndex.
+- [Workflows Introduction](https://docs.llamaindex.ai/en/stable/understanding/workflows/) - learn about LlamaIndex workflows.
+- [LlamaDeploy GitHub Repository](https://github.com/run-llama/llama_deploy)
+- [Chat-UI Documentation](https://ts.llamaindex.ai/docs/chat-ui)
+
+You can check out [the LlamaIndex GitHub repository](https://github.com/run-llama/llama_index) - your feedback and contributions are welcome!
@@ -0,0 +1,24 @@
+name: chat
+
+control-plane:
+  port: 8000
+
+default-service: workflow
+
+services:
+  workflow:
+    name: Workflow
+    source:
+      type: local
+      name: src
+    path: src/workflow:workflow
+    python-dependencies:
+      - llama-index-llms-openai>=0.4.5
+      - llama-index-core>=0.12.45
+
+ui:
+  name: My Nextjs App
+  port: 3000
+  source:
+    type: local
+    name: ui
@@ -0,0 +1,54 @@
+[project]
+name = "app"
+version = "0.1.0"
+description = ""
+readme = "README.md"
+requires-python = ">=3.11,<3.14"
+dependencies = [ "python-dotenv>=1.0.0,<2.0.0", "pydantic>=2.11.5", "aiostream>=0.5.2,<0.6.0", "llama-index-core>=0.12.28,<0.13.0", "llama-index-readers-file>=0.4.6,<1.0.0", "llama-index-indices-managed-llama-cloud>=0.6.3,<1.0.0", "llama-deploy", "docx2txt>=0.8,<0.9", "llama-index-llms-openai>=0.3.2,<0.4.0", "llama-index-embeddings-openai>=0.3.1,<0.4.0" ]
+
+[[project.authors]]
+name = "Marcus Schiesser"
+email = "mail@marcusschiesser.de"
+
+[project.optional-dependencies]
+dev = [ "mypy>=1.8.0,<2.0.0", "pytest>=8.3.5,<9.0.0", "pytest-asyncio>=0.25.3,<0.26.0" ]
+
+[project.scripts]
+generate = "src.generate:generate_index"
+
+[tool]
+[tool.uv]
+[tool.uv.sources]
+[tool.uv.sources.llama-deploy]
+git = "https://github.com/run-llama/llama_deploy"
+
+[tool.mypy]
+python_version = "3.11"
+plugins = "pydantic.mypy"
+exclude = [ "tests", "venv", ".venv", "output", "config" ]
+check_untyped_defs = true
+warn_unused_ignores = false
+show_error_codes = true
+namespace_packages = true
+ignore_missing_imports = true
+follow_imports = "silent"
+implicit_optional = true
+strict_optional = false
+disable_error_code = [ "return-value", "assignment" ]
+
+[[tool.mypy.overrides]]
+module = "src.*"
+ignore_missing_imports = false
+
+[tool.hatch]
+[tool.hatch.metadata]
+allow-direct-references = true
+
+[tool.hatch.build]
+[tool.hatch.build.targets]
+[tool.hatch.build.targets.wheel]
+packages = [ "src" ]
+
+[build-system]
+requires = [ "hatchling>=1.24" ]
+build-backend = "hatchling.build"
@@ -0,0 +1,106 @@
+from typing import Any, List, Optional
+
+from llama_index.core import QueryBundle
+from llama_index.core.postprocessor.types import BaseNodePostprocessor
+from llama_index.core.prompts import PromptTemplate
+from llama_index.core.query_engine.retriever_query_engine import RetrieverQueryEngine
+from llama_index.core.response_synthesizers import Accumulate
+from llama_index.core.schema import NodeWithScore
+from llama_index.core.tools.query_engine import QueryEngineTool
+
+
+# Used as a prompt for synthesizer
+# Override this prompt by setting the `CITATION_PROMPT` environment variable
+CITATION_PROMPT = """
+Context information is below.
+------------------
+{context_str}
+------------------
+The context are multiple text chunks, each text chunk has its own citation_id at the beginning.
+Use the citation_id for citation construction.
+
+Answer the following query with citations:
+------------------
+{query_str}
+------------------
+
+## Citation format
+
+[citation:id]
+
+Where:
+- [citation:] is a matching pattern which is required for all citations.
+- `id` is the `citation_id` provided in the context or previous response.
+
+Example:
+```
+    Here is a response that uses context information [citation:90ca859f-4f32-40ca-8cd0-edfad4fb298b] 
+    and other ideas that don't use context information [citation:17b2cc9a-27ae-4b6d-bede-5ca60fc00ff4] .\n
+    The citation block will be displayed automatically with useful information for the user in the UI [citation:1c606612-e75f-490e-8374-44e79f818d19] .
+```
+
+## Requirements:
+1. Always include citations for every fact from the context information in your response. 
+2. Make sure that the citation_id is correct with the context, don't mix up the citation_id with other information.
+
+Now, you answer the query with citations:
+"""
+
+
+class NodeCitationProcessor(BaseNodePostprocessor):
+    """
+    Add a new field `citation_id` to the metadata of the node by copying the id from the node.
+    Useful for citation construction.
+    """
+
+    def _postprocess_nodes(
+        self,
+        nodes: List[NodeWithScore],
+        query_bundle: Optional[QueryBundle] = None,
+    ) -> List[NodeWithScore]:
+        for node_score in nodes:
+            node_score.node.metadata["citation_id"] = node_score.node.node_id
+        return nodes
+
+
+class CitationSynthesizer(Accumulate):
+    """
+    Overload the Accumulate synthesizer to:
+    1. Update prepare node metadata for citation id
+    2. Update text_qa_template to include citations
+    """
+
+    def __init__(self, **kwargs: Any) -> None:
+        text_qa_template = kwargs.pop("text_qa_template", None)
+        if text_qa_template is None:
+            text_qa_template = PromptTemplate(template=CITATION_PROMPT)
+        super().__init__(text_qa_template=text_qa_template, **kwargs)
+
+
+# Add this prompt to your agent system prompt
+CITATION_SYSTEM_PROMPT = (
+    "\nAnswer the user question using the response from the query tool. "
+    "It's important to respect the citation information in the response. "
+    "Don't mix up the citation_id, keep them at the correct fact."
+)
+
+
+def enable_citation(query_engine_tool: QueryEngineTool) -> QueryEngineTool:
+    """
+    Enable citation for a query engine tool by using CitationSynthesizer and NodePostprocessor.
+    Note: This function will override the response synthesizer of your query engine.
+    """
+    query_engine = query_engine_tool.query_engine
+    if not isinstance(query_engine, RetrieverQueryEngine):
+        raise ValueError(
+            "Citation feature requires a RetrieverQueryEngine. Your tool's query engine is a "
+            f"{type(query_engine)}."
+        )
+    # Update the response synthesizer and node postprocessors
+    query_engine._response_synthesizer = CitationSynthesizer()
+    query_engine._node_postprocessors += [NodeCitationProcessor()]
+    query_engine_tool._query_engine = query_engine
+
+    # Update tool metadata
+    query_engine_tool.metadata.description += "\nThe output will include citations with the format [citation:id] for each chunk of information in the knowledge base."
+    return query_engine_tool
@@ -0,0 +1,37 @@
+import logging
+import os
+
+from dotenv import load_dotenv
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger()
+
+
+def generate_index():
+    """
+    Index the documents in the data directory.
+    """
+    from src.index import STORAGE_DIR
+    from src.settings import init_settings
+    from llama_index.core.indices import (
+        VectorStoreIndex,
+    )
+    from llama_index.core.readers import SimpleDirectoryReader
+
+    load_dotenv()
+    init_settings()
+
+    logger.info("Creating new index")
+    # load the documents and create the index
+    reader = SimpleDirectoryReader(
+        os.environ.get("DATA_DIR", "ui/data"),
+        recursive=True,
+    )
+    documents = reader.load_data()
+    index = VectorStoreIndex.from_documents(
+        documents,
+        show_progress=True,
+    )
+    # store it for later
+    index.storage_context.persist(STORAGE_DIR)
+    logger.info(f"Finished creating new index. Stored in {STORAGE_DIR}")
@@ -0,0 +1,21 @@
+import logging
+import os
+
+from llama_index.core.indices import load_index_from_storage
+from llama_index.core.storage import StorageContext
+
+logger = logging.getLogger("uvicorn")
+
+STORAGE_DIR = "src/storage"
+
+
+def get_index():
+    # check if storage already exists
+    if not os.path.exists(STORAGE_DIR):
+        return None
+    # load the existing index
+    logger.info(f"Loading index from {STORAGE_DIR}...")
+    storage_context = StorageContext.from_defaults(persist_dir=STORAGE_DIR)
+    index = load_index_from_storage(storage_context)
+    logger.info(f"Finished loading index from {STORAGE_DIR}")
+    return index
@@ -0,0 +1,47 @@
+import os
+from typing import Any, Optional
+
+from llama_index.core.base.base_query_engine import BaseQueryEngine
+from llama_index.core.indices.base import BaseIndex
+from llama_index.core.tools.query_engine import QueryEngineTool
+
+def create_query_engine(index: BaseIndex, **kwargs: Any) -> BaseQueryEngine:
+    """
+    Create a query engine for the given index.
+
+    Args:
+        index: The index to create a query engine for.
+        params (optional): Additional parameters for the query engine, e.g: similarity_top_k
+    """
+    top_k = int(os.getenv("TOP_K", 0))
+    if top_k != 0 and kwargs.get("filters") is None:
+        kwargs["similarity_top_k"] = top_k
+
+    return index.as_query_engine(**kwargs)
+
+
+def get_query_engine_tool(
+    index: BaseIndex,
+    name: Optional[str] = None,
+    description: Optional[str] = None,
+    **kwargs: Any,
+) -> QueryEngineTool:
+    """
+    Get a query engine tool for the given index.
+
+    Args:
+        index: The index to create a query engine for.
+        name (optional): The name of the tool.
+        description (optional): The description of the tool.
+    """
+    if name is None:
+        name = "query_index"
+    if description is None:
+        description = "Use this tool to retrieve information from a knowledge base. Provide a specific query and can call the tool multiple times if necessary."
+    query_engine = create_query_engine(index, **kwargs)
+    tool = QueryEngineTool.from_defaults(
+        query_engine=query_engine,
+        name=name,
+        description=description,
+    )
+    return tool
@@ -0,0 +1,14 @@
+import os
+
+from llama_index.core import Settings
+from llama_index.embeddings.openai import OpenAIEmbedding
+from llama_index.llms.openai import OpenAI
+
+
+def init_settings():
+    if os.getenv("OPENAI_API_KEY") is None:
+        raise RuntimeError("OPENAI_API_KEY is missing in environment variables")
+    Settings.llm = OpenAI(model=os.getenv("MODEL") or "gpt-4.1")
+    Settings.embed_model = OpenAIEmbedding(
+        model=os.getenv("EMBEDDING_MODEL") or "text-embedding-3-large"
+    )
@@ -0,0 +1,35 @@
+from dotenv import load_dotenv
+
+from llama_index.core.agent.workflow import AgentWorkflow
+from llama_index.core.settings import Settings
+
+from src.index import get_index
+from src.query import get_query_engine_tool
+from src.citation import CITATION_SYSTEM_PROMPT, enable_citation
+from src.settings import init_settings
+
+
+def create_workflow() -> AgentWorkflow:
+    load_dotenv()
+    init_settings()
+    index = get_index()
+    if index is None:
+        raise RuntimeError(
+            "Index not found! Please run `uv run generate` to index the data first."
+        )
+    # Create a query tool with citations enabled
+    query_tool = enable_citation(get_query_engine_tool(index=index))
+
+    # Define the system prompt for the agent
+    # Append the citation system prompt to the system prompt
+    system_prompt = """You are a helpful assistant"""
+    system_prompt += CITATION_SYSTEM_PROMPT
+
+    return AgentWorkflow.from_tools_or_functions(
+        tools_or_functions=[query_tool],
+        llm=Settings.llm,
+        system_prompt=system_prompt,
+    )
+
+
+workflow = create_workflow()
@@ -0,0 +1,10 @@
+import { LlamaIndexServer } from '@llamaindex/server'
+
+new LlamaIndexServer({
+  uiConfig: {
+    componentsDir: 'components',
+    layoutDir: 'layout',
+    llamaDeploy: { deployment: 'chat', workflow: 'workflow' },
+  },
+  port: 3000,
+}).start()
@@ -0,0 +1,40 @@
+'use client'
+
+import { Sparkles, Star } from 'lucide-react'
+
+export default function Header() {
+  return (
+    <div className="flex items-center justify-between p-2 px-4">
+      <div className="flex items-center gap-2">
+        <Sparkles className="size-4" />
+        <h1 className="font-semibold">LlamaIndex App</h1>
+      </div>
+      <div className="flex items-center justify-end gap-4">
+        <div className="flex items-center gap-2">
+          <a
+            href="https://www.llamaindex.ai/"
+            target="_blank"
+            rel="noopener noreferrer"
+            className="text-sm text-gray-600 hover:text-gray-800 dark:text-gray-400 dark:hover:text-gray-200"
+          >
+            Built by LlamaIndex
+          </a>
+          <img
+            className="h-[24px] w-[24px] rounded-sm"
+            src="https://ui.llamaindex.ai/llama.png"
+            alt="Llama Logo"
+          />
+        </div>
+        <a
+          href="https://github.com/run-llama/LlamaIndexTS"
+          target="_blank"
+          rel="noopener noreferrer"
+          className="hover:bg-accent flex items-center gap-2 rounded-md border border-gray-300 px-2 py-1 text-sm"
+        >
+          <Star className="size-4" />
+          Star on GitHub
+        </a>
+      </div>
+    </div>
+  )
+}
@@ -0,0 +1,18 @@
+{
+  "name": "llamaindex-server-ui",
+  "version": "0.0.1",
+  "private": true,
+  "scripts": {
+    "dev": "nodemon --exec tsx index.ts"
+  },
+  "dependencies": {
+    "@llamaindex/server": "latest",
+    "dotenv": "^16.4.7"
+  },
+  "devDependencies": {
+    "@types/node": "^20.10.3",
+    "nodemon": "^3.1.10",
+    "tsx": "4.7.2",
+    "typescript": "^5.3.2"
+  }
+}
@@ -5,8 +5,15 @@ import {
 } from '../../chat/annotations'
 import { JSONValue } from '../../chat/chat.interface'
 import { WorkflowEvent, WorkflowEventType } from '../use-workflow'
-import { AgentStreamEvent, SourceNodesEvent, UIEvent } from './types'
-import { SourceNode } from '../../widgets'
+import {
+  AgentStreamEvent,
+  RawNodeWithScore,
+  SourceNodesEvent,
+  ToolCallEvent,
+  ToolCallResultEvent,
+  UIEvent,
+} from './types'
+import { AgentEventData, SourceNode } from '../../widgets'

 /**
 * Transform a workflow event to message parts
@@ -16,7 +23,10 @@ import { SourceNode } from '../../widgets'
 * @param event - The event to transform
 * @returns The message parts (delta and annotations)
 */
-export function transformEventToMessageParts(event: WorkflowEvent): {
+export function transformEventToMessageParts(
+  event: WorkflowEvent,
+  fileServerUrl?: string
+): {
  delta: string
  annotations: MessageAnnotation<JSONValue>[]
 } {
@@ -31,7 +41,7 @@ export function transformEventToMessageParts(event: WorkflowEvent): {
    }
  }

-  const annotations = toVercelAnnotations(event)
+  const annotations = toVercelAnnotations(event, fileServerUrl)
  return { delta: '', annotations }
 }

@@ -51,7 +61,7 @@ function isInlineEvent(event: WorkflowEvent) {
  return inlineEventTypes.includes(event.type) && hasInlineData
 }

-function toVercelAnnotations(event: WorkflowEvent) {
+function toVercelAnnotations(event: WorkflowEvent, fileServerUrl?: string) {
  switch (event.type) {
    // convert source nodes event to source nodes annotation
    case WorkflowEventType.SourceNodesEvent.toString(): {
@@ -64,18 +74,14 @@ function toVercelAnnotations(event: WorkflowEvent) {
        return []
      }

-      const sources = nodes.map(({ node, score }) => ({
-        id: node.id_,
-        metadata: node.metadata,
-        score,
-        text: node.text,
-        url: (node.metadata?.URL as string) || '',
-      })) satisfies SourceNode[]
-
      return [
        {
          type: MessageAnnotationType.SOURCES,
-          data: { nodes: sources },
+          data: {
+            nodes: nodes.map(rawNode =>
+              convertRawNodeToSourceNode(rawNode, fileServerUrl)
+            ),
+          },
        },
      ]
    }
@@ -91,6 +97,72 @@ function toVercelAnnotations(event: WorkflowEvent) {
      ]
    }

+    // transfrom ToolCallEvent to AgentEvent
+    case WorkflowEventType.ToolCall.toString(): {
+      const { data } = event as ToolCallEvent
+
+      if (
+        'tool_name' in data &&
+        'tool_kwargs' in data &&
+        typeof data.tool_kwargs === 'object'
+      ) {
+        const { tool_name, tool_kwargs } = data
+
+        return [
+          {
+            type: MessageAnnotationType.AGENT_EVENTS,
+            data: {
+              agent: 'Agent',
+              text: `Calling tool: ${tool_name} with: ${JSON.stringify(tool_kwargs)}`,
+            } as AgentEventData,
+          },
+        ]
+      }
+
+      console.warn(
+        `Not supported this tool call event: ${JSON.stringify(event)}`
+      )
+
+      return []
+    }
+
+    // transform ToolCallResultEvent to suitable annotations
+    case WorkflowEventType.ToolCallResult.toString(): {
+      const { data } = event as ToolCallResultEvent
+      const { tool_output } = data
+
+      if (
+        'raw_output' in tool_output &&
+        typeof tool_output.raw_output === 'object' &&
+        tool_output.raw_output !== null
+      ) {
+        const { raw_output } = tool_output
+
+        // to source nodes annotation
+        if (
+          'source_nodes' in raw_output &&
+          Array.isArray(raw_output.source_nodes)
+        ) {
+          const rawNodes = raw_output.source_nodes as RawNodeWithScore[]
+          return [
+            {
+              type: MessageAnnotationType.SOURCES,
+              data: {
+                nodes: rawNodes.map(rawNode =>
+                  convertRawNodeToSourceNode(rawNode, fileServerUrl)
+                ),
+              },
+            },
+          ]
+        }
+      }
+
+      console.warn(
+        `Not supported this tool call result event: ${JSON.stringify(event)}`
+      )
+      return []
+    }
+
    // for other events which are not defined, convert them to vercel annotations with type is the qualified name of the event
    // eg: {"__is_pydantic": true, "value": {"item": "sample"}, "qualified_name": "myworkflow.MyEvent"}
    // will be converted to this annotation: {"type": "myworkflow.MyEvent", "data": {"item": "sample"}}
@@ -105,3 +177,54 @@ function toVercelAnnotations(event: WorkflowEvent) {
    }
  }
 }
+
+function convertRawNodeToSourceNode(
+  rawNode: RawNodeWithScore,
+  fileServerUrl?: string
+): SourceNode {
+  const { node, score } = rawNode
+
+  return {
+    id: node.id_,
+    metadata: node.metadata,
+    score,
+    text: node.text,
+    url: getDocumentUrlFromRawNode(rawNode, fileServerUrl),
+  }
+}
+
+function getDocumentUrlFromRawNode(
+  rawNode: RawNodeWithScore,
+  fileServerUrl?: string
+): string {
+  const { metadata } = rawNode.node
+  const { URL: fileURL, pipeline_id } = metadata
+
+  // if the file URL is provided in the metadata, use it
+  if (fileURL) return fileURL
+
+  // if `fileServerUrl` is not provided, return empty string
+  if (!fileServerUrl) {
+    console.warn(
+      `No 'fileServerUrl' provided. Nodes won't be displayed in ChatSources UI.`
+    )
+    return ''
+  }
+
+  let fileName = metadata?.file_name
+
+  // check if file_name is provided in the metadata
+  if (!fileName) {
+    console.warn(
+      `No file name found in this raw node: ${JSON.stringify(rawNode)}. It won't be displayed in ChatSources UI.`
+    )
+    return ''
+  }
+
+  if (pipeline_id) {
+    // pipeline_id is provided, the file is stored in LlamaCloud, make sure to cache it at the `fileServerUrl`
+    fileName = `${pipeline_id}$${fileName}`
+  }
+
+  return `${fileServerUrl}/${fileName}`
+}
@@ -15,6 +15,7 @@ export function useChatWorkflow({
  workflow,
  baseUrl,
  onError,
+  fileServerUrl,
 }: ChatWorkflowHookParams): ChatWorkflowHookHandler {
  const [input, setInput] = useState<string>('')
  const [messages, setMessages] = useState<Message[]>([])
@@ -51,7 +52,10 @@ export function useChatWorkflow({
    workflow,
    baseUrl,
    onData: event => {
-      const { delta, annotations } = transformEventToMessageParts(event)
+      const { delta, annotations } = transformEventToMessageParts(
+        event,
+        fileServerUrl
+      )
      updateLastMessage({ delta, annotations })
    },
  })
@@ -20,18 +20,30 @@ export interface AgentStreamEvent extends WorkflowEvent {
  }
 }

+// represent the raw node in the source nodes event
+export type RawNodeWithScore = {
+  node: {
+    id_: string
+    metadata: {
+      file_name: string | null
+      pipeline_id: string | null
+      pipeline_file_id: string | null
+      page_label: string | null
+      file_path: string | null
+      file_type: string | null
+      file_size: number | null
+      creation_date: string | null
+      last_modified_date: string | null
+      URL: string | null
+    }
+    text: string
+  }
+  score: number
+}
+
 export interface SourceNodesEvent extends WorkflowEvent {
  type: WorkflowEventType.SourceNodesEvent
-  data: {
-    nodes: {
-      node: {
-        id_: string
-        metadata: Record<string, JSONValue>
-        text: string
-      }
-      score: number
-    }[]
-  }
+  data: { nodes: RawNodeWithScore[] }
 }

 export interface UIEvent extends WorkflowEvent {
@@ -42,10 +54,37 @@ export interface UIEvent extends WorkflowEvent {
  }
 }

+export interface ToolCallEvent extends WorkflowEvent {
+  type: WorkflowEventType.ToolCall
+  data: {
+    tool_id: string
+    tool_name: string
+    tool_kwargs: JSONValue
+  }
+}
+
+export interface ToolCallResultEvent extends WorkflowEvent {
+  type: WorkflowEventType.ToolCallResult
+  data: {
+    tool_id: string
+    tool_name: string
+    tool_kwargs: JSONValue
+    tool_output: {
+      raw_output: JSONValue
+    }
+  }
+}
+
 export type ChatWorkflowHookParams = Pick<
  WorkflowHookParams,
  'deployment' | 'workflow' | 'baseUrl' | 'onError'
->
+> & {
+  // the endpoint to serve local files or llamacloud files
+  // if not provided, url for files cannot be constructed, and ChatSources won't be displayed in the UI
+  // eg. for non-llamacloud: localhost:3000/deployments/chat/ui/api/files/data
+  // eg. for llamacloud: localhost:3000/deployments/chat/ui/api/files/output/llamacloud
+  fileServerUrl?: string
+}

 export type ChatWorkflowHookHandler = ChatHandler & {
  resume: ChatWorkflowResume
@@ -100,13 +100,26 @@ export async function fetchTaskEvents<E extends WorkflowEvent>(
  try {
    callback?.onStart?.()

+    let retryParsedLines: string[] = []
+
    // eslint-disable-next-line no-constant-condition -- needed
    while (true) {
      const { done, value } = await reader.read()
      if (done) break

-      const chunk = decoder.decode(value, { stream: true })
-      const events = toWorkflowEvents<E>(chunk)
+      let chunk = decoder.decode(value, { stream: true })
+
+      if (retryParsedLines.length > 0) {
+        // if there are lines that failed to parse, append them to the current chunk
+        chunk = `${retryParsedLines.join('')}${chunk}`
+        console.info('Retry parsing with chunk:', { retryParsedLines, chunk })
+        retryParsedLines = [] // reset for next iteration
+      }
+
+      const { events, failedLines } = toWorkflowEvents<E>(chunk)
+
+      retryParsedLines.push(...failedLines)
+
      if (!events.length) continue

      accumulatedEvents.push(...events)
@@ -152,10 +165,15 @@ export async function sendEventToTask<E extends WorkflowEvent>(params: {
  return data.data
 }

-function toWorkflowEvents<E extends WorkflowEvent>(chunk: string): E[] {
+function toWorkflowEvents<E extends WorkflowEvent>(
+  chunk: string
+): {
+  events: E[]
+  failedLines: string[]
+} {
  if (typeof chunk !== 'string') {
    console.warn('Skipping non-string chunk:', chunk)
-    return []
+    return { events: [], failedLines: [] }
  }

  // One chunk can contain multiple events, so we need to parse each line
@@ -163,20 +181,47 @@ function toWorkflowEvents<E extends WorkflowEvent>(chunk: string): E[] {
    .trim()
    .split('\n')
    .filter(line => line.trim() !== '')
-  return lines.map(parseChunkLine<E>).filter(Boolean) as E[]
+
+  const parsedLines = lines.map(line => parseChunkLine<E>(line)).filter(Boolean)
+
+  // successfully parsed events
+  const events = parsedLines.map(line => line?.event).filter(Boolean) as E[]
+
+  // failed lines that could not be parsed into events
+  // will be merged into next chunk to re-try parsing
+  const failedLines = parsedLines
+    .filter(l => !l?.event)
+    .map(line => line?.line || '')
+    .filter(Boolean)
+
+  return {
+    events,
+    failedLines,
+  }
 }

-function parseChunkLine<E extends WorkflowEvent>(line: string): E | null {
+function parseChunkLine<E extends WorkflowEvent>(
+  line: string
+): {
+  line: string
+  event?: E | null
+} | null {
  try {
    const event = JSON.parse(line) as RawEvent
    if (!isRawEvent(event)) {
      console.warn('Skipping invalid raw event:', event)
      return null
    }
-    return { type: event.qualified_name, data: event.value } as E
+    return {
+      line,
+      event: { type: event.qualified_name, data: event.value } as E,
+    }
  } catch (error) {
    console.warn(`Failed to parse chunk in line: ${line}`, error)
-    return null
+    return {
+      line,
+      event: null,
+    }
  }
 }

@@ -42,6 +42,8 @@ export enum WorkflowEventType {
  SourceNodesEvent = 'llama_index.core.chat_ui.events.SourceNodesEvent',
  ArtifactEvent = 'llama_index.core.chat_ui.events.ArtifactEvent',
  UIEvent = 'llama_index.core.chat_ui.events.UIEvent',
+  ToolCall = 'llama_index.core.agent.workflow.workflow_events.ToolCall',
+  ToolCallResult = 'llama_index.core.agent.workflow.workflow_events.ToolCallResult',
 }

 export interface StreamingEventCallback<
@@ -355,4 +355,70 @@ describe('useWorkflow', () => {
      await waitFor(() => expect(result.current.events.length).toBeLessThan(13))
    })
  })
+
+  describe('Chunk parsing', () => {
+    it('should handle incomplete JSON chunks and retry parsing', async () => {
+      const originalHandlers = [...server.listHandlers()]
+
+      try {
+        const incompleteChunk1 = `{"__is_pydantic": true, "value": {"type": "ui_event", "data": {"id": null, "event": "analyze", "state": "inpro`
+        const incompleteChunk2 = `gress", "question": null, "answer": null}}, "qualified_name": "llama_index.core.chat_ui.events.UIEvent"}
+{"__is_pydantic": true, "value": {}, "qualified_name": "llama_index.core.workflow.events.StopEvent"}`
+
+        server.use(
+          http.get(
+            'http://127.0.0.1:4501/deployments/test-workflow/tasks/test-run-id/events',
+            () => {
+              const encoder = new TextEncoder()
+              const stream = new ReadableStream({
+                start(controller) {
+                  // Send incomplete chunk first
+                  controller.enqueue(encoder.encode(incompleteChunk1))
+
+                  // Wait a bit, then send the completion
+                  setTimeout(() => {
+                    controller.enqueue(encoder.encode(incompleteChunk2))
+                    controller.close()
+                  }, 10)
+                },
+              })
+
+              return new HttpResponse(stream, {
+                headers: { 'Content-Type': 'application/json' },
+              })
+            }
+          )
+        )
+
+        const { result } = renderHook(() =>
+          useWorkflow<TestEvent>(mockWorkflowParams)
+        )
+
+        await act(async () => {
+          await result.current.start({ message: 'test incomplete chunks' })
+        })
+
+        await waitFor(() => {
+          expect(result.current.events).toHaveLength(2)
+          expect(result.current.status).toBe('complete')
+        })
+
+        // First event should be the reconstructed incomplete chunk
+        const firstEvent = result.current.events[0]
+        expect(firstEvent.type).toBe('llama_index.core.chat_ui.events.UIEvent')
+        expect(firstEvent.data.type).toBe('ui_event')
+        expect(firstEvent.data.data.event).toBe('analyze')
+        expect(firstEvent.data.data.state).toBe('inprogress')
+        expect(firstEvent.data.data.id).toBeNull()
+
+        // Second event should be the stop event
+        const lastEvent = result.current.events[1]
+        expect(lastEvent.type).toBe(
+          'llama_index.core.workflow.events.StopEvent'
+        )
+      } finally {
+        server.resetHandlers(...originalHandlers)
+      }
+    })
+  })
 })
@@ -5,6 +5,4 @@ packages:
  - "docs"

  # UI will be built separately by llama-deploy server
-  - "!examples/llamadeploy/workflow/ui"
-  - "!examples/llamadeploy/chat/ui"
-  - "!examples/llamadeploy/with-tsserver/ui"
+  - "!examples/llamadeploy/*/ui"