add tool call events for agent_tool utils

.github/workflows/release_llama_index_server.yml

@@ -0,0 +1,122 @@
+name: Release llama-index-server
+
+on:
+  push:
+    branches:
+      - main
+    paths:
+      - "llama-index-server/**"
+      - ".github/workflows/release_llama_index_server.yml"
+  pull_request:
+    types:
+      - closed
+
+concurrency: ${{ github.workflow }}-${{ github.ref }}
+
+jobs:
+  release:
+    name: Create Release PR
+    runs-on: ubuntu-latest
+    defaults:
+      run:
+        working-directory: ./llama-index-server
+    if: github.event_name == 'push'
+
+    steps:
+      - name: Checkout Repository
+        uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+
+      - name: Install Poetry
+        run: |
+          curl -sSL https://install.python-poetry.org | python3 -
+
+      - name: Install dependencies
+        run: poetry install
+
+      - name: Setup Git
+        run: |
+          git config --global user.email "github-actions[bot]@users.noreply.github.com"
+          git config --global user.name "github-actions[bot]"
+
+      - name: Bump patch version
+        run: |
+          poetry version patch
+          git add pyproject.toml
+          git commit -m "chore(release): bump version to $(poetry version -s)"
+
+      - name: Get current version
+        id: get_version
+        run: |
+          version=$(poetry version -s)
+          echo "current_version=${version}" >> "$GITHUB_OUTPUT"
+
+      - name: Create Release PR
+        uses: peter-evans/create-pull-request@v6
+        with:
+          token: ${{ secrets.GITHUB_TOKEN }}
+          commit-message: "Release: llama-index-server v${{ steps.get_version.outputs.current_version }}"
+          title: "Release: llama-index-server v${{ steps.get_version.outputs.current_version }}"
+          body: |
+            This PR was automatically created to release a new version of the llama-index-server package.
+
+            Version: ${{ steps.get_version.outputs.current_version }}
+
+            Please review the changes and merge to trigger the release.
+          branch: release/llama-index-server-v${{ steps.get_version.outputs.current_version }}
+          base: main
+          labels: release, llama-index-server
+
+  publish:
+    name: Publish to PyPI
+    runs-on: ubuntu-latest
+    defaults:
+      run:
+        working-directory: ./llama-index-server
+    if: |
+      github.event_name == 'pull_request' && 
+      github.event.pull_request.merged == true && 
+      startsWith(github.event.pull_request.title, 'Release: llama-index-server') &&
+      startsWith(github.event.pull_request.head.ref, 'release/llama-index-server-v')
+
+    steps:
+      - name: Checkout Repository
+        uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+
+      - name: Install Poetry
+        run: |
+          curl -sSL https://install.python-poetry.org | python3 -
+
+      - name: Install dependencies
+        run: poetry install
+
+      - name: Build package
+        run: poetry build
+
+      - name: Publish to PyPI
+        env:
+          POETRY_PYPI_TOKEN_PYPI: ${{ secrets.LLAMA_INDEX_PYPI_TOKEN }}
+        run: poetry publish --build
+
+      - name: Create GitHub Release
+        uses: softprops/action-gh-release@v2
+        with:
+          tag_name: llama-index-server-v${{ steps.get_version.outputs.current_version }}
+          name: "llama-index-server v${{ steps.get_version.outputs.current_version }}"
+          body: |
+            Release of llama-index-server v${{ steps.get_version.outputs.current_version }}
+          draft: false
+          prerelease: false
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}

.github/workflows/test_llama_index_server.yml

@@ -0,0 +1,111 @@
+name: Build Package
+
+on:
+  pull_request:
+
+env:
+  POETRY_VERSION: "1.8.3"
+  PYTHON_VERSION: "3.9"
+
+jobs:
+  unit-test:
+    name: Unit Tests
+    runs-on: ${{ matrix.os }}
+    defaults:
+      run:
+        working-directory: llama-index-server
+    strategy:
+      matrix:
+        os: [ubuntu-latest, windows-latest]
+        python-version: ["3.9"]
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Install Poetry
+        run: pipx install poetry==${{ env.POETRY_VERSION }}
+
+      - name: Set up python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+          cache: "poetry"
+
+      - name: Configure Poetry
+        run: |
+          poetry config virtualenvs.create true
+          poetry config virtualenvs.in-project true
+          poetry env use python
+
+      - name: Install dependencies
+        shell: bash
+        run: poetry install --with dev
+
+      - name: Run unit tests
+        shell: bash
+        run: |
+          poetry run pytest tests
+
+  type-check:
+    name: Type Check
+    runs-on: ubuntu-latest
+    defaults:
+      run:
+        working-directory: llama-index-server
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Install Poetry
+        run: pipx install poetry==${{ env.POETRY_VERSION }}
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ env.PYTHON_VERSION }}
+          cache: "poetry"
+
+      - name: Configure Poetry
+        run: |
+          poetry config virtualenvs.create true
+          poetry config virtualenvs.in-project true
+          poetry env use python
+
+      - name: Install dependencies
+        shell: bash
+        run: poetry install --with dev
+
+      - name: Run mypy
+        shell: bash
+        run: poetry run mypy llama_index
+
+  build:
+    needs: [unit-test, type-check]
+    runs-on: ubuntu-latest
+    defaults:
+      run:
+        working-directory: llama-index-server
+    steps:
+      - uses: actions/checkout@v4
+      - name: Install Poetry
+        run: pipx install poetry==${{ env.POETRY_VERSION }}
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ env.PYTHON_VERSION }}
+      - name: Clear python cache
+        shell: bash
+        run: poetry cache clear --all pypi
+      - name: Build package
+        shell: bash
+        run: poetry build
+      - name: Test installing built package
+        shell: bash
+        run: python -m pip install .
+      - name: Test import
+        shell: bash
+        working-directory: ${{ vars.RUNNER_TEMP }}
+        run: python -c "from llama_index.server import LlamaIndexServer"
+      - name: Upload artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: llama-index-server
+          path: llama-index-server/dist/

.gitignore

@@ -8,6 +8,7 @@ node_modules
 
 # testing
 coverage
+.coverage
 
 # next.js
 .next/
@@ -48,6 +49,13 @@ e2e/cache
 
 # Python
 .mypy_cache/
+venv/
+.venv/
+dist/
+.__pycache__
+__pycache__
+.python-version
+.ui
 
 # build artifacts
 create-llama-*.tgz

llama-index-server/README.md

@@ -0,0 +1,55 @@
+# LlamaIndex Server
+
+LlamaIndexServer is a FastAPI application that allows you to quickly launch your workflow as an API server.
+
+## Installation
+
+```bash
+pip install llama-index-server
+```
+
+## Usage
+
+```python
+# main.py
+from llama_index.core.agent.workflow import AgentWorkflow
+from llama_index.core.workflow import Workflow
+from llama_index.core.tools import FunctionTool
+from llama_index.server import LlamaIndexServer
+
+
+# Define a factory function that returns a Workflow or AgentWorkflow
+def create_workflow() -> Workflow:
+    def fetch_weather(city: str) -> str:
+        return f"The weather in {city} is sunny"
+
+    return AgentWorkflow.from_tools(
+        tools=[
+            FunctionTool.from_defaults(
+                fn=fetch_weather,
+            )
+        ]
+    )
+
+
+# Create an API server the workflow
+app = LlamaIndexServer(
+    workflow_factory=create_workflow  # Supports Workflow or AgentWorkflow
+)
+```
+
+## Running the server
+
+- In the same directory as `main.py`, run the following command to start the server:
+
+  ```bash
+  fastapi dev
+  ```
+
+- Making a request to the server
+
+  ```bash
+  curl -X POST "http://localhost:8000/api/chat" -H "Content-Type: application/json" -d '{"message": "What is the weather in Tokyo?"}'
+  ```
+
+- See the API documentation at `http://localhost:8000/docs`

llama-index-server/llama_index/server/__init__.py

@@ -0,0 +1,3 @@
+from .server import LlamaIndexServer
+
+__all__ = ["LlamaIndexServer"]

llama-index-server/llama_index/server/api/callbacks/__init__.py

@@ -0,0 +1,11 @@
+from llama_index.server.api.callbacks.base import EventCallback
+from llama_index.server.api.callbacks.source_nodes import SourceNodesFromToolCall
+from llama_index.server.api.callbacks.suggest_next_questions import (
+    SuggestNextQuestions,
+)
+
+__all__ = [
+    "EventCallback",
+    "SourceNodesFromToolCall",
+    "SuggestNextQuestions",
+]

llama-index-server/llama_index/server/api/callbacks/base.py

@@ -0,0 +1,31 @@
+import logging
+from abc import ABC, abstractmethod
+from typing import Any
+
+logger = logging.getLogger("uvicorn")
+
+
+class EventCallback(ABC):
+    """
+    Base class for event callbacks during event streaming.
+    """
+
+    async def run(self, event: Any) -> Any:
+        """
+        Called for each event in the stream.
+        Default behavior: pass through the event unchanged.
+        """
+        return event
+
+    async def on_complete(self, final_response: str) -> Any:
+        """
+        Called when the stream is complete.
+        Default behavior: return None.
+        """
+        return None
+
+    @abstractmethod
+    def from_default(self, *args: Any, **kwargs: Any) -> "EventCallback":
+        """
+        Create a new instance of the processor from default values.
+        """

llama-index-server/llama_index/server/api/callbacks/source_nodes.py

@@ -0,0 +1,32 @@
+from typing import Any
+
+from llama_index.core.agent.workflow.workflow_events import ToolCallResult
+from llama_index.server.api.callbacks.base import EventCallback
+from llama_index.server.api.models import SourceNodesEvent
+
+
+class SourceNodesFromToolCall(EventCallback):
+    """
+    Extract source nodes from the query tool output.
+
+    Args:
+        query_tool_name: The name of the tool that queries the index.
+                         default is "query_index"
+    """
+
+    def __init__(self, query_tool_name: str = "query_index"):
+        self.query_tool_name = query_tool_name
+
+    def transform_tool_call_result(self, event: ToolCallResult) -> SourceNodesEvent:
+        source_nodes = event.tool_output.raw_output.source_nodes
+        return SourceNodesEvent(nodes=source_nodes)
+
+    async def run(self, event: Any) -> Any:
+        if isinstance(event, ToolCallResult):
+            if event.tool_name == self.query_tool_name:
+                return event, self.transform_tool_call_result(event)
+        return event
+
+    @classmethod
+    def from_default(cls, *args: Any, **kwargs: Any) -> "SourceNodesFromToolCall":
+        return cls()

llama-index-server/llama_index/server/api/callbacks/stream_handler.py

@@ -0,0 +1,69 @@
+import logging
+from typing import Any, AsyncGenerator, List, Optional
+
+from llama_index.core.workflow.handler import WorkflowHandler
+from llama_index.server.api.callbacks.base import EventCallback
+
+logger = logging.getLogger("uvicorn")
+
+
+class StreamHandler:
+    """
+    Streams events from a workflow handler through a chain of callbacks.
+    """
+
+    def __init__(
+        self,
+        workflow_handler: WorkflowHandler,
+        callbacks: Optional[List[EventCallback]] = None,
+    ):
+        self.workflow_handler = workflow_handler
+        self.callbacks = callbacks or []
+        self.accumulated_text = ""
+
+    async def cancel_run(self) -> None:
+        """Cancel the workflow handler."""
+        await self.workflow_handler.cancel_run()
+
+    async def stream_events(self) -> AsyncGenerator[Any, None]:
+        """Stream events through the processor chain."""
+        try:
+            async for event in self.workflow_handler.stream_events():
+                events_to_process = [event]
+                for callback in self.callbacks:
+                    next_events: list[Any] = []
+                    for evt in events_to_process:
+                        callback_output = await callback.run(evt)
+                        if isinstance(callback_output, (list, tuple)):
+                            next_events.extend(callback_output)
+                        elif callback_output is not None:
+                            next_events.append(callback_output)
+                    events_to_process = next_events
+
+                # Yield all processed events
+                for evt in events_to_process:
+                    yield evt
+
+            # After all events are processed, call on_complete for each callback
+            for callback in self.callbacks:
+                result = await callback.on_complete(self.accumulated_text)
+                if result:
+                    yield result
+
+        except Exception:
+            # Make sure to cancel the workflow on error
+            await self.workflow_handler.cancel_run()
+            raise
+
+    def accumulate_text(self, text: str) -> None:
+        """Accumulate text from the workflow handler."""
+        self.accumulated_text += text
+
+    @classmethod
+    def from_default(
+        cls,
+        handler: WorkflowHandler,
+        callbacks: Optional[List[EventCallback]] = None,
+    ) -> "StreamHandler":
+        """Create a new instance with the given workflow handler and callbacks."""
+        return cls(workflow_handler=handler, callbacks=callbacks)

llama-index-server/llama_index/server/api/callbacks/suggest_next_questions.py

@@ -0,0 +1,45 @@
+import logging
+from typing import Any, Optional
+
+from llama_index.server.api.callbacks.base import EventCallback
+from llama_index.server.api.models import ChatRequest
+from llama_index.server.services.suggest_next_question import (
+    SuggestNextQuestionsService,
+)
+
+logger = logging.getLogger("uvicorn")
+
+
+class SuggestNextQuestions(EventCallback):
+    """Processor for generating next question suggestions."""
+
+    def __init__(
+        self, chat_request: ChatRequest, logger: Optional[logging.Logger] = None
+    ):
+        self.chat_request = chat_request
+        self.accumulated_text = ""
+        if logger:
+            self.logger = logger
+        else:
+            self.logger = logging.getLogger("uvicorn")
+
+    async def on_complete(self, final_response: str) -> Any:
+        if final_response == "":
+            self.logger.warning(
+                "SuggestNextQuestions is enabled but final response is empty, make sure your content generator accumulates text"
+            )
+            return None
+
+        questions = await SuggestNextQuestionsService.run(
+            self.chat_request.messages, final_response
+        )
+        if questions:
+            return {
+                "type": "suggested_questions",
+                "data": questions,
+            }
+        return None
+
+    @classmethod
+    def from_default(cls, chat_request: ChatRequest) -> "SuggestNextQuestions":
+        return cls(chat_request=chat_request)

llama-index-server/llama_index/server/api/models.py

@@ -0,0 +1,137 @@
+import logging
+import os
+from enum import Enum
+from typing import Any, Dict, List, Optional
+
+from pydantic import BaseModel, Field, field_validator
+
+from llama_index.core.schema import NodeWithScore
+from llama_index.core.types import ChatMessage, MessageRole
+from llama_index.core.workflow import Event
+
+logger = logging.getLogger("uvicorn")
+
+
+class ChatConfig(BaseModel):
+    next_question_suggestions: bool = Field(
+        default=True,
+        description="Whether to suggest next questions",
+    )
+
+
+class ChatAPIMessage(BaseModel):
+    role: MessageRole
+    content: str
+
+    def to_llamaindex_message(self) -> ChatMessage:
+        return ChatMessage(role=self.role, content=self.content)
+
+
+class ChatRequest(BaseModel):
+    messages: List[ChatAPIMessage]
+    config: Optional[ChatConfig] = ChatConfig()
+
+    @field_validator("messages")
+    def validate_messages(cls, v: List[ChatAPIMessage]) -> List[ChatAPIMessage]:
+        if v[-1].role != MessageRole.USER:
+            raise ValueError("Last message must be from user")
+        return v
+
+
+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,
+            },
+        }
+
+
+class SourceNodesEvent(Event):
+    nodes: List[NodeWithScore]
+
+    def to_response(self) -> dict:
+        return {
+            "type": "sources",
+            "data": {
+                "nodes": [
+                    SourceNodes.from_source_node(node).model_dump()
+                    for node in self.nodes
+                ]
+            },
+        }
+
+
+class SourceNodes(BaseModel):
+    id: str
+    metadata: Dict[str, Any]
+    score: Optional[float]
+    text: str
+    url: Optional[str]
+
+    @classmethod
+    def from_source_node(cls, source_node: NodeWithScore) -> "SourceNodes":
+        metadata = source_node.node.metadata
+        url = cls.get_url_from_metadata(metadata)
+
+        return cls(
+            id=source_node.node.node_id,
+            metadata=metadata,
+            score=source_node.score,
+            text=source_node.node.text,  # type: ignore
+            url=url,
+        )
+
+    @classmethod
+    def get_url_from_metadata(
+        cls, metadata: Dict[str, Any], data_dir: Optional[str] = None
+    ) -> Optional[str]:
+        url_prefix = os.getenv("FILESERVER_URL_PREFIX")
+        if not url_prefix:
+            logger.warning(
+                "Warning: FILESERVER_URL_PREFIX not set in environment variables. Can't use file server"
+            )
+        if data_dir is None:
+            data_dir = "data"
+        file_name = metadata.get("file_name")
+
+        if file_name and url_prefix:
+            # file_name exists and file server is configured
+            pipeline_id = metadata.get("pipeline_id")
+            if pipeline_id:
+                # file is from LlamaCloud
+                file_name = f"{pipeline_id}${file_name}"
+                return f"{url_prefix}/output/llamacloud/{file_name}"
+            is_private = metadata.get("private", "false") == "true"
+            if is_private:
+                # file is a private upload
+                return f"{url_prefix}/output/uploaded/{file_name}"
+            # file is from calling the 'generate' script
+            # Get the relative path of file_path to data_dir
+            file_path = metadata.get("file_path")
+            data_dir = os.path.abspath(data_dir)
+            if file_path and data_dir:
+                relative_path = os.path.relpath(file_path, data_dir)
+                return f"{url_prefix}/data/{relative_path}"
+        # fallback to URL in metadata (e.g. for websites)
+        return metadata.get("URL")
+
+    @classmethod
+    def from_source_nodes(
+        cls, source_nodes: List[NodeWithScore]
+    ) -> List["SourceNodes"]:
+        return [cls.from_source_node(node) for node in source_nodes]

llama-index-server/llama_index/server/api/routers/chat.py

@@ -0,0 +1,109 @@
+import asyncio
+import logging
+from typing import AsyncGenerator, Callable, Union
+
+from fastapi import APIRouter, HTTPException
+from fastapi.responses import StreamingResponse
+
+from llama_index.core.agent.workflow.workflow_events import AgentStream
+from llama_index.core.workflow import StopEvent, Workflow
+from llama_index.server.api.callbacks import (
+    SourceNodesFromToolCall,
+    SuggestNextQuestions,
+)
+from llama_index.server.api.callbacks.base import EventCallback
+from llama_index.server.api.callbacks.stream_handler import StreamHandler
+from llama_index.server.api.models import ChatRequest
+from llama_index.server.api.utils.vercel_stream import VercelStreamResponse
+
+
+def chat_router(
+    workflow_factory: Callable[..., Workflow],
+    logger: logging.Logger,
+) -> APIRouter:
+    router = APIRouter(prefix="/chat")
+
+    @router.post("")
+    async def chat(request: ChatRequest) -> StreamingResponse:
+        try:
+            user_message = request.messages[-1].to_llamaindex_message()
+            chat_history = [
+                message.to_llamaindex_message() for message in request.messages[:-1]
+            ]
+            workflow = workflow_factory()
+            workflow_handler = workflow.run(
+                user_msg=user_message.content,
+                chat_history=chat_history,
+            )
+
+            callbacks: list[EventCallback] = [
+                SourceNodesFromToolCall(),
+            ]
+            if request.config and request.config.next_question_suggestions:
+                callbacks.append(SuggestNextQuestions(request))
+            stream_handler = StreamHandler(
+                workflow_handler=workflow_handler,
+                callbacks=callbacks,
+            )
+
+            return VercelStreamResponse(
+                content_generator=_stream_content(stream_handler, request, logger),
+            )
+        except Exception as e:
+            logger.error(e)
+            raise HTTPException(status_code=500, detail=str(e))
+
+    return router
+
+
+async def _stream_content(
+    handler: StreamHandler,
+    request: ChatRequest,
+    logger: logging.Logger,
+) -> AsyncGenerator[str, None]:
+    async def _text_stream(
+        event: Union[AgentStream, StopEvent],
+    ) -> AsyncGenerator[str, None]:
+        if isinstance(event, AgentStream):
+            if event.delta.strip():  # Only yield non-empty deltas
+                yield event.delta
+        elif isinstance(event, StopEvent):
+            if isinstance(event.result, str):
+                yield event.result
+            elif isinstance(event.result, AsyncGenerator):
+                async for chunk in event.result:
+                    if isinstance(chunk, str):
+                        yield chunk
+                    elif (
+                        hasattr(chunk, "delta") and chunk.delta.strip()
+                    ):  # Only yield non-empty deltas
+                        yield chunk.delta
+
+    stream_started = False
+    try:
+        async for event in handler.stream_events():
+            if not stream_started:
+                # Start the stream with an empty message
+                stream_started = True
+                yield VercelStreamResponse.convert_text("")
+
+            # Handle different types of events
+            if isinstance(event, (AgentStream, StopEvent)):
+                async for chunk in _text_stream(event):
+                    handler.accumulate_text(chunk)
+                    yield VercelStreamResponse.convert_text(chunk)
+            elif isinstance(event, dict):
+                yield VercelStreamResponse.convert_data(event)
+            elif hasattr(event, "to_response"):
+                event_response = event.to_response()
+                yield VercelStreamResponse.convert_data(event_response)
+            else:
+                yield VercelStreamResponse.convert_data(event.model_dump())
+
+    except asyncio.CancelledError:
+        logger.warning("Client cancelled the request!")
+        await handler.cancel_run()
+    except Exception as e:
+        logger.error(f"Error in stream response: {e}")
+        yield VercelStreamResponse.convert_error(str(e))
+        await handler.cancel_run()

llama-index-server/llama_index/server/api/utils/vercel_stream.py

@@ -0,0 +1,44 @@
+import json
+import logging
+from typing import Any, AsyncGenerator, Union
+
+from fastapi.responses import StreamingResponse
+
+logger = logging.getLogger("uvicorn")
+
+
+class VercelStreamResponse(StreamingResponse):
+    """
+    Converts preprocessed events into Vercel-compatible streaming response format.
+    """
+
+    TEXT_PREFIX = "0:"
+    DATA_PREFIX = "8:"
+    ERROR_PREFIX = "3:"
+
+    def __init__(
+        self,
+        content_generator: AsyncGenerator[str, None],
+        *args: Any,
+        **kwargs: Any,
+    ):
+        super().__init__(content_generator, *args, **kwargs)
+
+    @classmethod
+    def convert_text(cls, token: str) -> str:
+        """Convert text event to Vercel format."""
+        # Escape newlines and double quotes to avoid breaking the stream
+        token = json.dumps(token)
+        return f"{cls.TEXT_PREFIX}{token}\n"
+
+    @classmethod
+    def convert_data(cls, data: Union[dict, str]) -> str:
+        """Convert data event to Vercel format."""
+        data_str = json.dumps(data) if isinstance(data, dict) else data
+        return f"{cls.DATA_PREFIX}[{data_str}]\n"
+
+    @classmethod
+    def convert_error(cls, error: str) -> str:
+        """Convert error event to Vercel format."""
+        error_str = json.dumps(error)
+        return f"{cls.ERROR_PREFIX}{error_str}\n"

llama-index-server/llama_index/server/chat_ui.py

@@ -0,0 +1,55 @@
+import logging
+import shutil
+from pathlib import Path
+from typing import Optional
+
+import requests
+
+CHAT_UI_VERSION = "0.0.2"
+
+
+def download_chat_ui(
+    logger: Optional[logging.Logger] = None, target_path: str = ".ui"
+) -> None:
+    if logger is None:
+        logger = logging.getLogger("uvicorn")
+    path = Path(target_path)
+    temp_dir = _download_package(_get_download_link(CHAT_UI_VERSION))
+    _copy_ui_files(temp_dir, path)
+    logger.info("Chat UI downloaded and copied to static folder")
+
+
+def _get_download_link(version: str) -> str:
+    """Get the download link for the chat UI from the npm registry."""
+    return f"https://registry.npmjs.org/@llamaindex/server/-/server-{version}.tgz"
+
+
+def _download_package(url: str) -> Path:
+    """Download tar.gz file and extract all files into a temporary directory."""
+    import io
+    import tarfile
+    import tempfile
+
+    response = requests.get(url, headers={"User-Agent": "Mozilla/5.0"})
+    content = response.content
+
+    temp_dir = Path(tempfile.mkdtemp())
+
+    with tarfile.open(fileobj=io.BytesIO(content), mode="r:gz") as tar:
+        tar.extractall(path=temp_dir)
+
+    return temp_dir
+
+
+def _copy_ui_files(temp_dir: Path, target_path: Path) -> None:
+    """Copy files from the .next directory to the static directory."""
+    target_path.mkdir(parents=True, exist_ok=True)
+    next_dir = temp_dir / "package/dist/static"
+
+    if next_dir.exists():
+        for item in next_dir.iterdir():
+            dest = target_path / item.name
+            if item.is_dir():
+                shutil.copytree(item, dest, dirs_exist_ok=True)
+            else:
+                shutil.copy2(item, dest)

llama-index-server/llama_index/server/server.py

@@ -0,0 +1,134 @@
+import logging
+import os
+from typing import Any, Callable, Optional
+
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.staticfiles import StaticFiles
+
+from llama_index.core.workflow import Workflow
+from llama_index.server.api.routers.chat import chat_router
+from llama_index.server.chat_ui import download_chat_ui
+
+
+class LlamaIndexServer(FastAPI):
+    workflow_factory: Callable[..., Workflow]
+    api_prefix: str = "/api"
+    include_ui: Optional[bool]
+    verbose: bool = False
+    ui_path: str = ".ui"
+
+    def __init__(
+        self,
+        workflow_factory: Callable[..., Workflow],
+        logger: Optional[logging.Logger] = None,
+        use_default_routers: Optional[bool] = False,
+        env: Optional[str] = None,
+        include_ui: Optional[bool] = None,
+        verbose: bool = False,
+        *args: Any,
+        **kwargs: Any,
+    ):
+        """
+        Initialize the LlamaIndexServer.
+
+        Args:
+            workflow_factory: A factory function that creates a workflow instance for each request.
+            logger: The logger to use.
+            use_default_routers: Whether to use the default routers (chat, mount `data` and `output` directories).
+            env: The environment to run the server in.
+            include_ui: Whether to show an chat UI in the root path.
+            verbose: Whether to show verbose logs.
+        """
+        super().__init__(*args, **kwargs)
+
+        self.workflow_factory = workflow_factory
+        self.logger = logger or logging.getLogger("uvicorn")
+        self.verbose = verbose
+        self.include_ui = include_ui  # Store the explicitly passed value first
+
+        if use_default_routers:
+            self.add_default_routers()
+
+        if str(env).lower() == "dev":
+            self.allow_cors("*")
+            if self.include_ui is None:
+                self.include_ui = True
+        if self.include_ui is None:
+            self.include_ui = False
+
+        if self.include_ui:
+            self.mount_ui()
+
+    # Default routers
+    def add_default_routers(self) -> None:
+        self.add_chat_router()
+        self.mount_data_dir()
+        self.mount_output_dir()
+
+    def add_chat_router(self) -> None:
+        """
+        Add the chat router.
+        """
+        self.include_router(
+            chat_router(
+                self.workflow_factory,
+                self.logger,
+            ),
+            prefix=self.api_prefix,
+        )
+
+    def mount_ui(self) -> None:
+        """
+        Mount the UI.
+        """
+        # Check if the static folder exists
+        if self.include_ui:
+            if not os.path.exists(self.ui_path):
+                self.logger.warning(
+                    f"UI files not found, downloading UI to {self.ui_path}"
+                )
+                download_chat_ui(logger=self.logger, target_path=self.ui_path)
+            self._mount_static_files(directory=self.ui_path, path="/", html=True)
+
+    def mount_data_dir(self, data_dir: str = "data") -> None:
+        """
+        Mount the data directory.
+        """
+        self._mount_static_files(
+            directory=data_dir, path=f"{self.api_prefix}/files/data", html=True
+        )
+
+    def mount_output_dir(self, output_dir: str = "output") -> None:
+        """
+        Mount the output directory.
+        """
+        self._mount_static_files(
+            directory=output_dir, path=f"{self.api_prefix}/files/output", html=True
+        )
+
+    def _mount_static_files(
+        self, directory: str, path: str, html: bool = False
+    ) -> None:
+        """
+        Mount static files from a directory if it exists.
+        """
+        if os.path.exists(directory):
+            self.logger.info(f"Mounting static files '{directory}' at '{path}'")
+            self.mount(
+                path,
+                StaticFiles(directory=directory, check_dir=False, html=html),
+                name=f"{directory}-static",
+            )
+
+    def allow_cors(self, origin: str = "*") -> None:
+        """
+        Allow CORS for a specific origin.
+        """
+        self.add_middleware(
+            CORSMiddleware,
+            allow_origins=[origin],
+            allow_credentials=True,
+            allow_methods=["*"],
+            allow_headers=["*"],
+        )

llama-index-server/llama_index/server/services/file.py

@@ -0,0 +1,117 @@
+import logging
+import os
+import re
+import uuid
+from pathlib import Path
+from typing import List, Optional, Union
+
+from pydantic import BaseModel, Field
+
+logger = logging.getLogger(__name__)
+
+PRIVATE_STORE_PATH = str(Path("output", "uploaded"))
+TOOL_STORE_PATH = str(Path("output", "tools"))
+LLAMA_CLOUD_STORE_PATH = str(Path("output", "llamacloud"))
+
+
+class DocumentFile(BaseModel):
+    id: str
+    name: str  # Stored file name
+    type: Optional[str] = None
+    size: Optional[int] = None
+    url: Optional[str] = None
+    path: Optional[str] = Field(
+        None,
+        description="The stored file path. Used internally in the server.",
+        exclude=True,
+    )
+    refs: Optional[List[str]] = Field(
+        None, description="The document ids in the index."
+    )
+
+
+class FileService:
+    """
+    To store the files uploaded by the user.
+    """
+
+    @classmethod
+    def save_file(
+        cls,
+        content: Union[bytes, str],
+        file_name: str,
+        save_dir: Optional[str] = None,
+    ) -> DocumentFile:
+        """
+        Save the content to a file in the local file server (accessible via URL).
+
+        Args:
+            content (bytes | str): The content to save, either bytes or string.
+            file_name (str): The original name of the file.
+            save_dir (Optional[str]): The relative path from the current working directory. Defaults to the `output/uploaded` directory.
+
+        Returns:
+            The metadata of the saved file.
+        """
+        if save_dir is None:
+            save_dir = os.path.join("output", "uploaded")
+
+        file_id = str(uuid.uuid4())
+        name, extension = os.path.splitext(file_name)
+        extension = extension.lstrip(".")
+        sanitized_name = _sanitize_file_name(name)
+        if extension == "":
+            raise ValueError("File is not supported!")
+        new_file_name = f"{sanitized_name}_{file_id}.{extension}"
+
+        file_path = os.path.join(save_dir, new_file_name)
+
+        if isinstance(content, str):
+            content = content.encode()
+
+        try:
+            os.makedirs(os.path.dirname(file_path), exist_ok=True)
+            with open(file_path, "wb") as file:
+                file.write(content)
+        except PermissionError as e:
+            logger.error(f"Permission denied when writing to file {file_path}: {e!s}")
+            raise
+        except OSError as e:
+            logger.error(f"IO error occurred when writing to file {file_path}: {e!s}")
+            raise
+        except Exception as e:
+            logger.error(f"Unexpected error when writing to file {file_path}: {e!s}")
+            raise
+
+        logger.info(f"Saved file to {file_path}")
+
+        file_url_prefix = os.getenv("FILESERVER_URL_PREFIX")
+        if file_url_prefix is None:
+            logger.warning(
+                "FILESERVER_URL_PREFIX is not set. Some features may not work correctly."
+            )
+            file_url_prefix = "http://localhost:8000/api/files"
+        file_size = os.path.getsize(file_path)
+
+        file_url = os.path.join(
+            file_url_prefix,
+            save_dir,
+            new_file_name,
+        )
+
+        return DocumentFile(
+            id=file_id,
+            name=new_file_name,
+            type=extension,
+            size=file_size,
+            path=file_path,
+            url=file_url,
+            refs=None,
+        )
+
+
+def _sanitize_file_name(file_name: str) -> str:
+    """
+    Sanitize the file name by replacing all non-alphanumeric characters with underscores.
+    """
+    return re.sub(r"[^a-zA-Z0-9.]", "_", file_name)

llama-index-server/llama_index/server/services/suggest_next_question.py

@@ -0,0 +1,95 @@
+import logging
+import os
+import re
+from typing import List, Optional, Union
+
+from llama_index.core.prompts import PromptTemplate
+from llama_index.core.settings import Settings
+from llama_index.server.api.models import ChatAPIMessage
+
+logger = logging.getLogger("uvicorn")
+
+
+class SuggestNextQuestionsService:
+    """
+    Suggest the next questions that user might ask based on the conversation history.
+    """
+
+    prompt = PromptTemplate(
+        r"""
+You're a helpful assistant! Your task is to suggest the next questions that user might interested in to keep the conversation going.
+Here is the conversation history
+---------------------
+{conversation}
+---------------------
+Given the conversation history, please give me 3 questions that user might ask next!
+Your answer should be wrapped in three sticks without any index numbers and follows the following format:
+\`\`\`
+<question 1>
+<question 2>
+<question 3>
+\`\`\`
+"""
+    )
+
+    @classmethod
+    def get_configured_prompt(cls) -> PromptTemplate:
+        prompt = os.getenv("NEXT_QUESTION_PROMPT", None)
+        if not prompt:
+            return cls.prompt
+        return PromptTemplate(prompt)
+
+    @classmethod
+    async def suggest_next_questions_all_messages(
+        cls,
+        messages: List[ChatAPIMessage],
+    ) -> Optional[List[str]]:
+        """
+        Suggest the next questions that user might ask based on the conversation history.
+        """
+        prompt_template = cls.get_configured_prompt()
+
+        try:
+            # Reduce the cost by only using the last two messages
+            last_user_message = None
+            last_assistant_message = None
+            for message in reversed(messages):
+                if message.role == "user":
+                    last_user_message = f"User: {message.content}"
+                elif message.role == "assistant":
+                    last_assistant_message = f"Assistant: {message.content}"
+                if last_user_message and last_assistant_message:
+                    break
+            conversation: str = f"{last_user_message}\n{last_assistant_message}"
+
+            # Call the LLM and parse questions from the output
+            prompt = prompt_template.format(conversation=conversation)
+            output = await Settings.llm.acomplete(prompt)
+            return cls._extract_questions(output.text)
+
+        except Exception as e:
+            logger.error(f"Error when generating next question: {e}")
+            return None
+
+    @classmethod
+    def _extract_questions(cls, text: str) -> Union[List[str], None]:
+        content_match = re.search(r"```(.*?)```", text, re.DOTALL)
+        content = content_match.group(1) if content_match else None
+        if not content:
+            return None
+        return [q.strip() for q in content.split("\n") if q.strip()]
+
+    @classmethod
+    async def run(
+        cls,
+        chat_history: List[ChatAPIMessage],
+        response: str,
+    ) -> Optional[List[str]]:
+        """
+        Suggest the next questions that user might ask based on the chat history and the last response.
+        """
+        messages = [
+            *chat_history,
+            ChatAPIMessage(role="assistant", content=response),  # type: ignore
+        ]
+        return await cls.suggest_next_questions_all_messages(messages)

llama-index-server/llama_index/server/tools/document_generator.py

@@ -0,0 +1,236 @@
+import logging
+import os
+import re
+from enum import Enum
+from io import BytesIO
+
+from llama_index.core.tools.function_tool import FunctionTool
+
+OUTPUT_DIR = "output/tools"
+
+
+class DocumentType(Enum):
+    PDF = "pdf"
+    HTML = "html"
+
+
+COMMON_STYLES = """
+body {
+    font-family: Arial, sans-serif;
+    line-height: 1.3;
+    color: #333;
+}
+h1, h2, h3, h4, h5, h6 {
+    margin-top: 1em;
+    margin-bottom: 0.5em;
+}
+p {
+    margin-bottom: 0.7em;
+}
+code {
+    background-color: #f4f4f4;
+    padding: 2px 4px;
+    border-radius: 4px;
+}
+pre {
+    background-color: #f4f4f4;
+    padding: 10px;
+    border-radius: 4px;
+    overflow-x: auto;
+}
+table {
+    border-collapse: collapse;
+    width: 100%;
+    margin-bottom: 1em;
+}
+th, td {
+    border: 1px solid #ddd;
+    padding: 8px;
+    text-align: left;
+}
+th {
+    background-color: #f2f2f2;
+    font-weight: bold;
+}
+"""
+
+HTML_SPECIFIC_STYLES = """
+body {
+    max-width: 800px;
+    margin: 0 auto;
+    padding: 20px;
+}
+"""
+
+PDF_SPECIFIC_STYLES = """
+@page {
+    size: letter;
+    margin: 2cm;
+}
+body {
+    font-size: 11pt;
+}
+h1 { font-size: 18pt; }
+h2 { font-size: 16pt; }
+h3 { font-size: 14pt; }
+h4, h5, h6 { font-size: 12pt; }
+pre, code {
+    font-family: Courier, monospace;
+    font-size: 0.9em;
+}
+"""
+
+HTML_TEMPLATE = """
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <style>
+        {common_styles}
+        {specific_styles}
+    </style>
+</head>
+<body>
+    {content}
+</body>
+</html>
+"""
+
+
+class DocumentGenerator:
+    @classmethod
+    def _generate_html_content(cls, original_content: str) -> str:
+        """
+        Generate HTML content from the original markdown content.
+        """
+        try:
+            import markdown  # type: ignore
+        except ImportError:
+            raise ImportError(
+                "Failed to import required modules. Please install markdown."
+            )
+
+        # Convert markdown to HTML with fenced code and table extensions
+        return markdown.markdown(original_content, extensions=["fenced_code", "tables"])
+
+    @classmethod
+    def _generate_pdf(cls, html_content: str) -> BytesIO:
+        """
+        Generate a PDF from the HTML content.
+        """
+        try:
+            from xhtml2pdf import pisa
+        except ImportError:
+            raise ImportError(
+                "Failed to import required modules. Please install xhtml2pdf."
+            )
+
+        pdf_html = HTML_TEMPLATE.format(
+            common_styles=COMMON_STYLES,
+            specific_styles=PDF_SPECIFIC_STYLES,
+            content=html_content,
+        )
+
+        buffer = BytesIO()
+        pdf = pisa.pisaDocument(
+            BytesIO(pdf_html.encode("UTF-8")), buffer, encoding="UTF-8"
+        )
+
+        if pdf.err:
+            logging.error(f"PDF generation failed: {pdf.err}")
+            raise ValueError("PDF generation failed")
+
+        buffer.seek(0)
+        return buffer
+
+    @classmethod
+    def _generate_html(cls, html_content: str) -> str:
+        """
+        Generate a complete HTML document with the given HTML content.
+        """
+        return HTML_TEMPLATE.format(
+            common_styles=COMMON_STYLES,
+            specific_styles=HTML_SPECIFIC_STYLES,
+            content=html_content,
+        )
+
+    @classmethod
+    def generate_document(
+        cls, original_content: str, document_type: str, file_name: str
+    ) -> str:
+        """
+        To generate document as PDF or HTML file.
+        Parameters:
+            original_content: str (markdown style)
+            document_type: str (pdf or html) specify the type of the file format based on the use case
+            file_name: str (name of the document file) must be a valid file name, no extensions needed
+        Returns:
+            str (URL to the document file): A file URL ready to serve.
+        """
+        try:
+            doc_type = DocumentType(document_type.lower())
+        except ValueError:
+            raise ValueError(
+                f"Invalid document type: {document_type}. Must be 'pdf' or 'html'."
+            )
+        # Always generate html content first
+        html_content = cls._generate_html_content(original_content)
+
+        # Based on the type of document, generate the corresponding file
+        if doc_type == DocumentType.PDF:
+            content = cls._generate_pdf(html_content)
+            file_extension = "pdf"
+        elif doc_type == DocumentType.HTML:
+            content = BytesIO(cls._generate_html(html_content).encode("utf-8"))
+            file_extension = "html"
+        else:
+            raise ValueError(f"Unexpected document type: {document_type}")
+
+        file_name = cls._validate_file_name(file_name)
+        file_path = os.path.join(OUTPUT_DIR, f"{file_name}.{file_extension}")
+
+        cls._write_to_file(content, file_path)
+
+        return f"{os.getenv('FILESERVER_URL_PREFIX')}/{OUTPUT_DIR}/{file_name}.{file_extension}"
+
+    @staticmethod
+    def _write_to_file(content: BytesIO, file_path: str) -> None:
+        """
+        Write the content to a file.
+        """
+        try:
+            os.makedirs(os.path.dirname(file_path), exist_ok=True)
+            with open(file_path, "wb") as file:
+                file.write(content.getvalue())
+        except Exception:
+            raise
+
+    @staticmethod
+    def _validate_file_name(file_name: str) -> str:
+        """
+        Validate the file name.
+        """
+        # Don't allow directory traversal
+        if os.path.isabs(file_name):
+            raise ValueError("File name is not allowed.")
+        # Don't allow special characters
+        if re.match(r"^[a-zA-Z0-9_.-]+$", file_name):
+            return file_name
+        else:
+            raise ValueError("File name is not allowed to contain special characters.")
+
+    @classmethod
+    def _validate_packages(cls) -> None:
+        try:
+            import markdown  # noqa: F401
+            import xhtml2pdf  # noqa: F401
+        except ImportError:
+            raise ImportError(
+                "Failed to import required modules. Please install markdown and xhtml2pdf "
+                "using `pip install markdown xhtml2pdf`"
+            )
+
+    def to_tool(self) -> FunctionTool:
+        self._validate_packages()
+        return FunctionTool.from_defaults(self.generate_document)

llama-index-server/llama_index/server/tools/index/__init__.py

@@ -0,0 +1,3 @@
+from .query import get_query_engine_tool
+
+__all__ = ["get_query_engine_tool"]

llama-index-server/llama_index/server/tools/index/query.py

@@ -0,0 +1,49 @@
+import os
+from typing import Any, Optional
+
+from llama_index.core.base.base_query_engine import BaseQueryEngine
+from llama_index.core.tools.query_engine import QueryEngineTool
+from llama_index.core.indices.base import BaseIndex
+
+
+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 about the text corpus from an index."
+        )
+    query_engine = create_query_engine(index, **kwargs)
+    return QueryEngineTool.from_defaults(
+        query_engine=query_engine,
+        name=name,
+        description=description,
+    )

llama-index-server/llama_index/server/tools/index/utils.py

@@ -0,0 +1,13 @@
+from datetime import timedelta
+
+from cachetools import TTLCache, cached  # type: ignore
+
+from llama_index.core.storage import StorageContext
+
+
+@cached(
+    TTLCache(maxsize=10, ttl=timedelta(minutes=5).total_seconds()),
+    key=lambda *args, **kwargs: "global_storage_context",
+)
+def get_storage_context(persist_dir: str) -> StorageContext:
+    return StorageContext.from_defaults(persist_dir=persist_dir)

llama-index-server/llama_index/server/tools/interpreter.py

@@ -0,0 +1,227 @@
+import base64
+import logging
+import os
+import uuid
+from typing import Any, List, Optional
+
+from pydantic import BaseModel
+
+from llama_index.core.tools import FunctionTool
+from llama_index.server.services.file import DocumentFile, FileService
+
+logger = logging.getLogger("uvicorn")
+
+
+class InterpreterExtraResult(BaseModel):
+    type: str
+    content: Optional[str] = None
+    filename: Optional[str] = None
+    url: Optional[str] = None
+
+
+class E2BToolOutput(BaseModel):
+    is_error: bool
+    logs: "Logs"  # type: ignore # noqa: F821
+    error_message: Optional[str] = None
+    results: List[InterpreterExtraResult] = []
+    retry_count: int = 0
+
+
+class E2BCodeInterpreter:
+    output_dir = "output/tools"
+    uploaded_files_dir = "output/uploaded"
+    interpreter: Optional["Sandbox"] = None  # type: ignore # noqa: F821
+
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        filesever_url_prefix: Optional[str] = None,
+        output_dir: Optional[str] = None,
+        uploaded_files_dir: Optional[str] = None,
+    ):
+        """
+        Args:
+            api_key: The API key for the E2B Code Interpreter. If not provided, it will be read from the environment variable `E2B_API_KEY`.
+            filesever_url_prefix: The prefix for the file server or loaded from env: `FILESERVER_URL_PREFIX`, default is `/api/files`.
+            output_dir: The directory for the output files. Default is `output/tools`.
+            uploaded_files_dir: The directory for the files to be uploaded to the sandbox. Default is `output/uploaded`.
+        """
+        self._validate_package()
+        if api_key is None:
+            api_key = os.getenv("E2B_API_KEY")
+        if filesever_url_prefix is None:
+            filesever_url_prefix = os.getenv("FILESERVER_URL_PREFIX", "/api/files")
+        if not api_key:
+            raise ValueError(
+                "E2B_API_KEY key is required to run code interpreter. Get it here: https://e2b.dev/docs/getting-started/api-key"
+            )
+        if output_dir is not None:
+            self.output_dir = output_dir
+        if uploaded_files_dir is not None:
+            self.uploaded_files_dir = uploaded_files_dir
+        self.filesever_url_prefix = filesever_url_prefix
+        self.interpreter = None
+        self.api_key = api_key
+
+    @classmethod
+    def _validate_package(cls) -> None:
+        try:
+            from e2b_code_interpreter import Sandbox  # noqa: F401
+            from e2b_code_interpreter.models import Logs  # noqa: F401
+        except ImportError:
+            raise ImportError(
+                "e2b_code_interpreter is not installed. Please install it using `pip install e2b-code-interpreter`."
+            )
+
+    def __del__(self) -> None:
+        """
+        Kill the interpreter when the tool is no longer in use.
+        """
+        if self.interpreter is not None:
+            self.interpreter.kill()
+
+    def _init_interpreter(self, sandbox_files: List[str] = []) -> None:
+        """
+        Lazily initialize the interpreter.
+        """
+        from e2b_code_interpreter import Sandbox
+
+        logger.info(f"Initializing interpreter with {len(sandbox_files)} files")
+        self.interpreter = Sandbox(api_key=self.api_key)
+        if len(sandbox_files) > 0:
+            for file_path in sandbox_files:
+                file_name = os.path.basename(file_path)
+                local_file_path = os.path.join(self.uploaded_files_dir, file_name)
+                with open(local_file_path, "rb") as f:
+                    content = f.read()
+                    if self.interpreter and self.interpreter.files:
+                        self.interpreter.files.write(file_path, content)
+            logger.info(f"Uploaded {len(sandbox_files)} files to sandbox")
+
+    def _save_to_disk(self, base64_data: str, ext: str) -> DocumentFile:
+        buffer = base64.b64decode(base64_data)
+
+        # Output from e2b doesn't have a name. Create a random name for it.
+        filename = f"e2b_file_{uuid.uuid4()}.{ext}"
+
+        return FileService.save_file(
+            buffer, file_name=filename, save_dir=self.output_dir
+        )
+
+    def _parse_result(self, result: Any) -> List[InterpreterExtraResult]:
+        """
+        The result could include multiple formats (e.g. png, svg, etc.) but encoded in base64
+        We save each result to disk and return saved file metadata (extension, filename, url).
+        """
+        if not result:
+            return []
+
+        output = []
+
+        try:
+            formats = result.formats()
+            results = [result[format] for format in formats]
+
+            for ext, data in zip(formats, results):
+                if ext in ["png", "svg", "jpeg", "pdf"]:
+                    document_file = self._save_to_disk(data, ext)
+                    output.append(
+                        InterpreterExtraResult(
+                            type=ext,
+                            filename=document_file.name,
+                            url=document_file.url,
+                        )
+                    )
+                else:
+                    # Try serialize data to string
+                    try:
+                        data = str(data)
+                    except Exception as e:
+                        data = f"Error when serializing data: {e}"
+                    output.append(
+                        InterpreterExtraResult(
+                            type=ext,
+                            content=data,
+                        )
+                    )
+        except Exception as error:
+            logger.exception(error, exc_info=True)
+            logger.error("Error when parsing output from E2b interpreter tool", error)
+
+        return output
+
+    def interpret(
+        self,
+        code: str,
+        sandbox_files: List[str] = [],
+        retry_count: int = 0,
+    ) -> E2BToolOutput:
+        """
+        Execute Python code in a Jupyter notebook cell. The tool will return the result, stdout, stderr, display_data, and error.
+        If the code needs to use a file, ALWAYS pass the file path in the sandbox_files argument.
+        You have a maximum of 3 retries to get the code to run successfully.
+
+        Parameters:
+            code (str): The Python code to be executed in a single cell.
+            sandbox_files (List[str]): List of local file paths to be used by the code. The tool will throw an error if a file is not found.
+            retry_count (int): Number of times the tool has been retried.
+        """
+        from e2b_code_interpreter.models import Logs
+
+        if retry_count > 2:
+            return E2BToolOutput(
+                is_error=True,
+                logs=Logs(
+                    stdout="",
+                    stderr="",
+                    display_data="",
+                    error="",
+                ),
+                error_message="Failed to execute the code after 3 retries. Explain the error to the user and suggest a fix.",
+                retry_count=retry_count,
+            )
+
+        if self.interpreter is None:
+            self._init_interpreter(sandbox_files)
+
+        if self.interpreter:
+            logger.info(
+                f"\n{'=' * 50}\n> Running following AI-generated code:\n{code}\n{'=' * 50}"
+            )
+            exec = self.interpreter.run_code(code)
+
+            if exec.error:
+                error_message = f"The code failed to execute successfully. Error: {exec.error}. Try to fix the code and run again."
+                logger.error(error_message)
+                # Calling the generated code caused an error. Kill the interpreter and return the error to the LLM so it can try to fix the error
+                try:
+                    self.interpreter.kill()  # type: ignore
+                except Exception:
+                    pass
+                finally:
+                    self.interpreter = None
+                output = E2BToolOutput(
+                    is_error=True,
+                    logs=exec.logs,
+                    results=[],
+                    error_message=error_message,
+                    retry_count=retry_count + 1,
+                )
+            else:
+                if len(exec.results) == 0:
+                    output = E2BToolOutput(is_error=False, logs=exec.logs, results=[])
+                else:
+                    results = self._parse_result(exec.results[0])
+                    output = E2BToolOutput(
+                        is_error=False,
+                        logs=exec.logs,
+                        results=results,
+                        retry_count=retry_count + 1,
+                    )
+            return output
+        else:
+            raise ValueError("Interpreter is not initialized.")
+
+    def to_tool(self) -> FunctionTool:
+        self._validate_package()
+        return FunctionTool.from_defaults(self.interpret)

llama-index-server/llama_index/server/utils/agent_tool.py

@@ -0,0 +1,255 @@
+import logging
+import uuid
+from abc import ABC, abstractmethod
+from typing import Any, AsyncGenerator, Callable, Optional
+
+from pydantic import BaseModel, ConfigDict
+
+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 llama_index.server.api.models import AgentRunEvent, AgentRunEventType
+from llama_index.core.agent.workflow.workflow_events import ToolCall, ToolCallResult
+
+logger = logging.getLogger("uvicorn")
+
+class ToolCallOutput(BaseModel):
+    tool_call_id: str
+    tool_output: ToolOutput
+
+
+class ContextAwareTool(FunctionTool, ABC):
+    @abstractmethod
+    async def acall(self, ctx: Context, input: Any) -> ToolOutput:  # type: ignore
+        pass
+
+
+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 or []}
+        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:
+        if not self.has_tool_calls():
+            raise ValueError("No tool calls")
+        if self.is_calling_different_tools():
+            raise ValueError("Calling different tools")
+        return self.tool_calls[0].tool_name  # type: ignore
+
+    async def full_response(self) -> str:
+        assert self.generator is not None
+        full_response = ""
+        async for chunk in self.generator:
+            content = chunk.delta  # type: ignore
+            if content:
+                full_response += 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,  # type: ignore
+        )
+
+
+async def call_tools(
+    ctx: Context,
+    agent_name: str,
+    tools: list[BaseTool],
+    tool_calls: list[ToolSelection],
+    emit_agent_events: bool = True,
+) -> list[ToolCallOutput]:
+    """
+    Call tools and return the tool call responses.
+    """
+    if len(tool_calls) == 0:
+        return []
+    tools_by_name = {tool.metadata.get_name(): tool for tool in tools}
+    if len(tool_calls) == 1:
+        if emit_agent_events:
+            ctx.write_event_to_stream(
+                AgentRunEvent(
+                    name=agent_name,
+                    msg=f"{tool_calls[0].tool_name}: {tool_calls[0].tool_kwargs}",
+                )
+            )
+        return [
+            await call_tool(
+                ctx,
+                tools_by_name[tool_calls[0].tool_name],
+                tool_calls[0]
+            )
+        ]
+    # Multiple tool calls, show progress
+    tool_call_outputs: list[ToolCallOutput] = []
+
+    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_call_outputs.append(
+                ToolCallOutput(
+                    tool_call_id=tool_call.tool_id,
+                    tool_output=ToolOutput(
+                        is_error=True,
+                        content=f"Tool {tool_call.tool_name} does not exist",
+                        tool_name=tool_call.tool_name,
+                        raw_input=tool_call.tool_kwargs,
+                        raw_output={
+                            "error": f"Tool {tool_call.tool_name} does not exist",
+                        },
+                    )
+                )
+            )
+            continue
+
+        tool_call_output = await call_tool(
+            ctx,
+            tool,
+            tool_call,
+        )
+        if emit_agent_events:
+            ctx.write_event_to_stream(
+                AgentRunEvent(
+                    name=agent_name,
+                    msg=f"{tool_call.tool_name}: {tool_call.tool_kwargs}",
+                    event_type=AgentRunEventType.PROGRESS,
+                    data={
+                        "id": progress_id,
+                        "total": total_steps,
+                        "current": i,
+                    },
+                )
+            )
+        tool_call_outputs.append(tool_call_output)
+    return tool_call_outputs
+
+
+async def call_tool(
+    ctx: Context,
+    tool: BaseTool,
+    tool_call: ToolSelection,
+) -> ToolCallOutput:
+    ctx.write_event_to_stream(
+        ToolCall(
+            tool_name=tool_call.tool_name,
+            tool_id=tool_call.tool_id,
+            tool_kwargs=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
+            output = await tool.acall(ctx=ctx, **tool_call.tool_kwargs)
+        else:
+            output = await tool.acall(**tool_call.tool_kwargs)  # type: ignore
+    except Exception as e:
+        logger.error(f"Got error in tool {tool_call.tool_name}: {e!s}")
+        output = ToolOutput(
+            is_error=True,
+            content=f"Error: {e!s}",
+            tool_name=tool.metadata.get_name(),
+            raw_input=tool_call.tool_kwargs,
+            raw_output={
+                "error": str(e),
+            },
+        )
+    ctx.write_event_to_stream(
+        ToolCallResult(
+            tool_name=tool_call.tool_name,
+            tool_kwargs=tool_call.tool_kwargs,
+            tool_id=tool_call.tool_id,
+            tool_output=output,
+            return_direct=False,
+        )
+    )
+    return ToolCallOutput(
+        tool_call_id=tool_call.tool_id,
+        tool_output=output,
+    )
+
+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

llama-index-server/pyproject.toml

@@ -0,0 +1,61 @@
+[build-system]
+build-backend = "poetry.core.masonry.api"
+requires = ["poetry-core"]
+
+[tool.codespell]
+check-filenames = true
+check-hidden = true
+# Feel free to un-skip examples, and experimental, you will just need to
+# work through many typos (--write-changes and --interactive will help)
+skip = "*.csv,*.html,*.json,*.jsonl,*.pdf,*.txt,*.ipynb"
+
+[tool.mypy]
+disallow_untyped_defs = true
+# Remove venv skip when integrated with pre-commit
+exclude = ["_static", "build", "examples", "notebooks", "venv"]
+ignore_missing_imports = true
+namespace_packages = true
+explicit_package_bases = true
+python_version = "3.10"
+
+[tool.poetry]
+authors = ["Your Name <you@example.com>"]
+description = "llama-index fastapi server"
+exclude = ["**/BUILD"]
+license = "MIT"
+name = "llama-index-server"
+packages = [{include = "llama_index/"}]
+readme = "README.md"
+version = "0.1.0"
+
+[tool.poetry.dependencies]
+python = ">=3.9,<4.0"
+fastapi = {extras = ["standard"], version = "^0.115.11"}
+cachetools = "^5.5.2"
+requests = "^2.32.3"
+llama-index-core = "^0.12.0"
+llama-index-readers-file = "^0.4.6"
+
+[tool.poetry.group.dev.dependencies]
+black = {extras = ["jupyter"], version = "<=23.9.1,>=23.7.0"}
+codespell = {extras = ["toml"], version = ">=v2.2.6"}
+e2b-code-interpreter = "^1.1.1"
+ipython = "8.10.0"
+jupyter = "^1.0.0"
+markdown = "^3.7"
+mypy = "1.15.0"
+pre-commit = "3.2.0"
+pylint = "2.15.10"
+pytest = "^8.3.5"
+pytest-asyncio = "^0.25.3"
+pytest-mock = "3.11.1"
+ruff = "0.0.292"
+tree-sitter-languages = "^1.8.0"
+types-Deprecated = ">=0.1.0"
+types-PyYAML = "^6.0.12.12"
+types-protobuf = "^4.24.0.4"
+types-redis = "4.5.5.0"
+types-requests = "2.28.11.8"  # TODO: unpin when mypy>0.991
+types-setuptools = "67.1.0.0"
+xhtml2pdf = "^0.2.17"
+pytest-cov = "^6.0.0"

llama-index-server/tests/api/test_chat_api.py

@@ -0,0 +1,149 @@
+import logging
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+from fastapi import FastAPI
+from httpx import ASGITransport, AsyncClient
+
+from llama_index.core.workflow import StopEvent, Workflow
+from llama_index.core.workflow.handler import WorkflowHandler
+from llama_index.server.api.models import ChatAPIMessage, ChatRequest
+from llama_index.server.api.routers.chat import chat_router
+
+
+@pytest.fixture()
+def logger():
+    return logging.getLogger("test")
+
+
+@pytest.fixture()
+def chat_request():
+    """Create a simple chat request with one user message."""
+    return ChatRequest(
+        messages=[ChatAPIMessage(role="user", content="Hello, how are you?")]
+    )
+
+
+@pytest.fixture()
+def mock_workflow():
+    """Create a mock workflow that returns a simple response."""
+    workflow = MagicMock(spec=Workflow)
+    handler = AsyncMock(spec=WorkflowHandler)
+
+    # Setup the handler to stream a simple response event
+    async def mock_stream_events():
+        yield StopEvent(result="I'm doing well, thank you for asking!")
+
+    handler.stream_events.return_value = mock_stream_events()
+    workflow.run.return_value = handler
+
+    return workflow
+
+
+@pytest.fixture()
+def workflow_factory(mock_workflow):
+    """Create a factory function that returns our mock workflow."""
+
+    def factory(verbose=False):
+        return mock_workflow
+
+    return factory
+
+
+@pytest.mark.asyncio()
+async def test_chat_router(chat_request, workflow_factory, logger):
+    """Test that the chat router handles a request correctly."""
+    # Create a FastAPI app and mount our router
+    app = FastAPI()
+    router = chat_router(workflow_factory, logger)
+    app.include_router(router)
+
+    # Make a request to the chat endpoint
+    async with AsyncClient(
+        transport=ASGITransport(app=app), base_url="http://test"
+    ) as client:
+        response = await client.post("/chat", json=chat_request.model_dump())
+
+        # Check response status
+        assert response.status_code == 200
+
+        # For streaming responses we don't check the content-type header directly
+        # Instead, check that we get the expected content in the response body
+
+        # The response is a stream, so we need to collect the chunks
+        content = response.content.decode()
+
+        # Verify content structure follows expected format
+        assert "0:" in content  # Text prefix for VercelStreamResponse
+        # Verify if the response contains the expected message
+        assert "I'm doing well" in content
+
+        # Verify the mock workflow was called correctly
+        mock_workflow = workflow_factory()
+        mock_workflow.run.assert_called_once()
+
+        # Verify the workflow was called with the correct arguments
+        call_args = mock_workflow.run.call_args[1]
+        assert call_args["user_msg"] == "Hello, how are you?"
+        assert isinstance(call_args["chat_history"], list)
+        assert len(call_args["chat_history"]) == 0  # No history for first message
+
+
+@pytest.mark.asyncio()
+async def test_chat_with_agent_workflow(logger):
+    """Test that the chat router works with a workflow that mimics an agent workflow."""
+    # Create a simple workflow that mimics an agent workflow
+    mock_workflow = MagicMock(spec=Workflow)
+    handler = AsyncMock(spec=WorkflowHandler)
+
+    # Setup the handler to stream a simple response about weather
+    async def mock_stream_events():
+        yield StopEvent(
+            result="The weather in New York is sunny. I used the weather tool to get this information."
+        )
+
+    handler.stream_events.return_value = mock_stream_events()
+    mock_workflow.run.return_value = handler
+
+    # Create a factory function that returns our mock workflow
+    def workflow_factory(verbose=False):
+        return mock_workflow
+
+    # Create a FastAPI app and mount our router
+    app = FastAPI()
+    router = chat_router(workflow_factory, logger)
+    app.include_router(router)
+
+    # Create a chat request asking about weather
+    chat_request = ChatRequest(
+        messages=[
+            ChatAPIMessage(role="user", content="What's the weather in New York?")
+        ]
+    )
+
+    # Make a request to the chat endpoint
+    async with AsyncClient(
+        transport=ASGITransport(app=app), base_url="http://test"
+    ) as client:
+        response = await client.post("/chat", json=chat_request.model_dump())
+
+        # Check response status
+        assert response.status_code == 200
+
+        # The response is a stream, so we need to collect the chunks
+        content = response.content.decode()
+
+        # Verify content structure follows expected format
+        assert "0:" in content  # Text prefix for VercelStreamResponse
+
+        # Verify the response content contains expected keywords
+        assert "weather" in content and "New York" in content and "sunny" in content
+
+        # Verify the mock workflow was called correctly
+        mock_workflow.run.assert_called_once()
+
+        # Verify the workflow was called with the correct arguments
+        call_args = mock_workflow.run.call_args[1]
+        assert call_args["user_msg"] == "What's the weather in New York?"
+        assert isinstance(call_args["chat_history"], list)
+        assert len(call_args["chat_history"]) == 0  # No history for first message

llama-index-server/tests/api/test_event_stream.py

@@ -0,0 +1,250 @@
+import asyncio
+import logging
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from llama_index.core.agent.workflow.workflow_events import AgentStream
+from llama_index.core.workflow import StopEvent
+from llama_index.core.workflow.handler import WorkflowHandler
+from llama_index.server.api.models import ChatAPIMessage, ChatRequest
+from llama_index.server.api.routers.chat import _stream_content
+from llama_index.server.api.utils.vercel_stream import VercelStreamResponse
+
+
+@pytest.fixture()
+def logger():
+    return logging.getLogger("test")
+
+
+@pytest.fixture()
+def chat_request():
+    return ChatRequest(messages=[ChatAPIMessage(role="user", content="test message")])
+
+
+@pytest.fixture()
+def mock_workflow_handler():
+    handler = AsyncMock(spec=WorkflowHandler)
+    handler.accumulate_text = MagicMock()
+    return handler
+
+
+class TestEventStream:
+    @pytest.mark.asyncio()
+    async def test_stream_content_with_agent_stream(
+        self, mock_workflow_handler, chat_request, logger
+    ):
+        # Setup
+        mock_workflow_handler.stream_events.return_value = (
+            self._mock_agent_stream_events()
+        )
+
+        # Execute
+        result = [
+            chunk
+            async for chunk in _stream_content(
+                mock_workflow_handler, chat_request, logger
+            )
+        ]
+
+        # Assert
+        assert len(result) == 3  # Empty start + 2 text chunks
+        assert result[0] == VercelStreamResponse.convert_text("")
+        assert result[1] == VercelStreamResponse.convert_text("Hello")
+        assert result[2] == VercelStreamResponse.convert_text(" World")
+
+    @pytest.mark.asyncio()
+    async def test_stream_content_with_stop_event_string(
+        self, mock_workflow_handler, chat_request, logger
+    ):
+        # Setup
+        mock_workflow_handler.stream_events.return_value = (
+            self._mock_stop_event_string()
+        )
+
+        # Execute
+        result = [
+            chunk
+            async for chunk in _stream_content(
+                mock_workflow_handler, chat_request, logger
+            )
+        ]
+
+        # Assert
+        assert len(result) == 2  # Empty start + result string
+        assert result[0] == VercelStreamResponse.convert_text("")
+        assert result[1] == VercelStreamResponse.convert_text("Final answer")
+
+    @pytest.mark.asyncio()
+    async def test_stream_content_with_stop_event_delta_objects(
+        self, mock_workflow_handler, chat_request, logger
+    ):
+        # Setup
+        mock_workflow_handler.stream_events.return_value = (
+            self._mock_stop_event_delta_objects()
+        )
+
+        # Execute
+        result = [
+            chunk
+            async for chunk in _stream_content(
+                mock_workflow_handler, chat_request, logger
+            )
+        ]
+
+        # Assert
+        assert len(result) == 3  # Empty start + 2 delta chunks
+        assert result[0] == VercelStreamResponse.convert_text("")
+        assert result[1] == VercelStreamResponse.convert_text("Delta 1")
+        assert result[2] == VercelStreamResponse.convert_text("Delta 2")
+
+    @pytest.mark.asyncio()
+    async def test_stream_content_with_event_with_to_response(
+        self, mock_workflow_handler, chat_request, logger
+    ):
+        # Setup
+        mock_workflow_handler.stream_events.return_value = (
+            self._mock_event_with_to_response()
+        )
+
+        # Execute
+        result = [
+            chunk
+            async for chunk in _stream_content(
+                mock_workflow_handler, chat_request, logger
+            )
+        ]
+
+        # Assert
+        assert len(result) == 2  # Empty start + event with to_response
+        assert result[0] == VercelStreamResponse.convert_text("")
+        assert result[1] == VercelStreamResponse.convert_data({"event_type": "test"})
+
+    @pytest.mark.asyncio()
+    async def test_stream_content_with_event_with_model_dump(
+        self, mock_workflow_handler, chat_request, logger
+    ):
+        # Setup
+        mock_workflow_handler.stream_events.return_value = (
+            self._mock_event_with_model_dump()
+        )
+
+        # Execute
+        result = [
+            chunk
+            async for chunk in _stream_content(
+                mock_workflow_handler, chat_request, logger
+            )
+        ]
+
+        # Assert
+        assert len(result) == 2  # Empty start + event with model_dump
+        assert result[0] == VercelStreamResponse.convert_text("")
+        assert result[1] == VercelStreamResponse.convert_data(None)
+
+    @pytest.mark.asyncio()
+    async def test_stream_content_with_cancelled_error(
+        self, mock_workflow_handler, chat_request, logger
+    ):
+        # Setup
+        mock_workflow_handler.stream_events.side_effect = asyncio.CancelledError()
+        logger.warning = MagicMock()
+
+        # Execute
+        result = [
+            chunk
+            async for chunk in _stream_content(
+                mock_workflow_handler, chat_request, logger
+            )
+        ]
+
+        # Assert
+        assert len(result) == 0
+        mock_workflow_handler.cancel_run.assert_called_once()
+        logger.warning.assert_called_once()
+
+    @pytest.mark.asyncio()
+    async def test_stream_content_with_exception(
+        self, mock_workflow_handler, chat_request, logger
+    ):
+        # Setup
+        error_message = "Test error"
+        mock_workflow_handler.stream_events.side_effect = Exception(error_message)
+        logger.error = MagicMock()
+
+        # Execute
+        result = [
+            chunk
+            async for chunk in _stream_content(
+                mock_workflow_handler, chat_request, logger
+            )
+        ]
+
+        # Assert
+        assert len(result) == 1
+        assert result[0] == VercelStreamResponse.convert_error(error_message)
+        mock_workflow_handler.cancel_run.assert_called_once()
+        logger.error.assert_called_once()
+
+    async def _mock_agent_stream_events(self):
+        yield AgentStream(
+            delta="Hello", response="", current_agent_name="", tool_calls=[], raw=""
+        )
+        yield AgentStream(
+            delta=" World", response="", current_agent_name="", tool_calls=[], raw=""
+        )
+
+    async def _mock_agent_stream_with_empty_deltas(self):
+        yield AgentStream(
+            delta="   ",  # Empty delta with spaces - should be filtered
+            response="",
+            current_agent_name="",
+            tool_calls=[],
+            raw="",
+        )
+        yield AgentStream(
+            delta="Valid delta",
+            response="",
+            current_agent_name="",
+            tool_calls=[],
+            raw="",
+        )
+        yield AgentStream(
+            delta="\n",  # Newline-only delta - should be filtered
+            response="",
+            current_agent_name="",
+            tool_calls=[],
+            raw="",
+        )
+
+    async def _mock_stop_event_string(self):
+        yield StopEvent(result="Final answer")
+
+    async def _mock_stop_event_delta_objects(self):
+        async def generator():
+            # Create proper objects with delta attribute that can be serialized
+            class ObjectWithDelta:
+                def __init__(self, delta_value) -> None:
+                    self.delta = delta_value
+
+            yield ObjectWithDelta("Delta 1")
+            yield ObjectWithDelta("Delta 2")
+            yield ObjectWithDelta("  ")  # Should be filtered out by strip check
+
+        yield StopEvent(result=generator())
+
+    async def _mock_dict_event(self):
+        yield {"key": "value"}
+
+    async def _mock_event_with_to_response(self):
+        event = MagicMock()
+        event.to_response.return_value = {"event_type": "test"}
+        yield event
+
+    async def _mock_event_with_model_dump(self):
+        event = MagicMock()
+        event.model_dump.return_value = {"name": "test_event"}
+        # Override to_response to return None - this means convert_data(None) will be called
+        event.to_response = MagicMock(return_value=None)
+        # The model_dump value is ignored when to_response returns None
+        yield event

llama-index-server/tests/services/test_file_service.py

@@ -0,0 +1,207 @@
+import os
+import uuid
+from unittest.mock import mock_open, patch
+
+import pytest
+
+from llama_index.server.services.file import FileService, _sanitize_file_name
+
+
+class TestFileService:
+    def test_sanitize_file_name(self):
+        # Test with normal alphanumeric name
+        assert _sanitize_file_name("test123") == "test123"
+
+        # Test with spaces
+        assert _sanitize_file_name("test file") == "test_file"
+
+        # Test with special characters
+        assert _sanitize_file_name("test@file!name") == "test_file_name"
+
+        # Test with path-like characters
+        assert _sanitize_file_name("test/file/name") == "test_file_name"
+
+        # Test with dots (should be preserved)
+        assert _sanitize_file_name("test.file.name") == "test.file.name"
+
+    @patch("uuid.uuid4")
+    @patch("os.path.getsize")
+    @patch("builtins.open", new_callable=mock_open)
+    @patch("os.makedirs")
+    def test_save_file_string_content(
+        self, mock_makedirs, mock_file_open, mock_getsize, mock_uuid
+    ):
+        # Setup
+        test_uuid = "12345678-1234-5678-1234-567812345678"
+        mock_uuid.return_value = uuid.UUID(test_uuid)
+        mock_getsize.return_value = 11  # Length of "Hello World"
+
+        # Execute
+        result = FileService.save_file(
+            content="Hello World", file_name="test.txt", save_dir="test_dir"
+        )
+
+        # Assert
+        expected_path = os.path.join("test_dir", f"test_{test_uuid}.txt")
+        mock_makedirs.assert_called_once_with(
+            os.path.dirname(expected_path), exist_ok=True
+        )
+        mock_file_open.assert_called_once_with(expected_path, "wb")
+        mock_file_open().write.assert_called_once_with(b"Hello World")
+
+        assert result.id == test_uuid
+        assert result.name == f"test_{test_uuid}.txt"
+        assert result.type == "txt"
+        assert result.size == 11
+        assert result.path == expected_path
+        assert result.url.endswith(expected_path)
+        assert result.refs is None
+
+    @patch("uuid.uuid4")
+    @patch("os.path.getsize")
+    @patch("builtins.open", new_callable=mock_open)
+    @patch("os.makedirs")
+    def test_save_file_bytes_content(
+        self, mock_makedirs, mock_file_open, mock_getsize, mock_uuid
+    ):
+        # Setup
+        test_uuid = "12345678-1234-5678-1234-567812345678"
+        mock_uuid.return_value = uuid.UUID(test_uuid)
+        mock_getsize.return_value = 11  # Length of "Hello World"
+
+        # Execute
+        result = FileService.save_file(
+            content=b"Hello World", file_name="test.txt", save_dir="test_dir"
+        )
+
+        # Assert
+        expected_path = os.path.join("test_dir", f"test_{test_uuid}.txt")
+        mock_makedirs.assert_called_once_with(
+            os.path.dirname(expected_path), exist_ok=True
+        )
+        mock_file_open.assert_called_once_with(expected_path, "wb")
+        mock_file_open().write.assert_called_once_with(b"Hello World")
+        assert result.path == expected_path
+        assert result.type == "txt"
+
+    @patch("uuid.uuid4")
+    @patch("os.path.getsize")
+    @patch("builtins.open", new_callable=mock_open)
+    @patch("os.makedirs")
+    def test_save_file_with_special_characters(
+        self, mock_makedirs, mock_file_open, mock_getsize, mock_uuid
+    ):
+        # Setup
+        test_uuid = "12345678-1234-5678-1234-567812345678"
+        mock_uuid.return_value = uuid.UUID(test_uuid)
+        mock_getsize.return_value = 11
+
+        # Execute
+        result = FileService.save_file(
+            content="Hello World", file_name="test@file!.txt", save_dir="test_dir"
+        )
+
+        # Assert
+        expected_path = os.path.join("test_dir", f"test_file__{test_uuid}.txt")
+        mock_makedirs.assert_called_once_with(
+            os.path.dirname(expected_path), exist_ok=True
+        )
+        mock_file_open.assert_called_once_with(expected_path, "wb")
+        assert result.path == expected_path
+        assert result.name == f"test_file__{test_uuid}.txt"
+
+    @patch("uuid.uuid4")
+    @patch("os.path.getsize")
+    @patch("builtins.open", new_callable=mock_open)
+    @patch("os.makedirs")
+    def test_save_file_default_directory(
+        self, mock_makedirs, mock_file_open, mock_getsize, mock_uuid
+    ):
+        # Setup
+        test_uuid = "12345678-1234-5678-1234-567812345678"
+        mock_uuid.return_value = uuid.UUID(test_uuid)
+        mock_getsize.return_value = 11
+
+        # Execute
+        result = FileService.save_file(content="Hello World", file_name="test.txt")
+
+        # Assert
+        expected_path = os.path.join("output", "uploaded", f"test_{test_uuid}.txt")
+        mock_makedirs.assert_called_once_with(
+            os.path.dirname(expected_path), exist_ok=True
+        )
+        assert result.path == expected_path
+
+    @patch("uuid.uuid4")
+    @patch("os.getenv")
+    @patch("os.path.getsize")
+    @patch("builtins.open", new_callable=mock_open)
+    @patch("os.makedirs")
+    def test_save_file_custom_url_prefix(
+        self, mock_makedirs, mock_file_open, mock_getsize, mock_getenv, mock_uuid
+    ):
+        # Setup
+        test_uuid = "12345678-1234-5678-1234-567812345678"
+        mock_uuid.return_value = uuid.UUID(test_uuid)
+        mock_getsize.return_value = 11
+        mock_getenv.return_value = "https://custom-url.com/files"
+
+        # Execute
+        result = FileService.save_file(
+            content="Hello World", file_name="test.txt", save_dir="test_dir"
+        )
+
+        # Assert
+        expected_path = os.path.join("test_dir", f"test_{test_uuid}.txt")
+        mock_makedirs.assert_called_once_with(
+            os.path.dirname(expected_path), exist_ok=True
+        )
+        mock_file_open.assert_called_once_with(expected_path, "wb")
+        assert result.path == expected_path
+        expected_url = os.path.join(
+            "https://custom-url.com/files", "test_dir", f"test_{test_uuid}.txt"
+        )
+        assert result.url == expected_url
+
+    def test_save_file_no_extension(self):
+        # Test that saving a file without extension raises ValueError
+        with pytest.raises(ValueError, match="File is not supported!"):
+            FileService.save_file(
+                content="Hello World", file_name="test", save_dir="test_dir"
+            )
+
+    @patch("uuid.uuid4")
+    @patch("os.path.getsize")
+    @patch("builtins.open")
+    @patch("os.makedirs")
+    def test_save_file_permission_error(
+        self, mock_makedirs, mock_file_open, mock_getsize, mock_uuid
+    ):
+        # Setup
+        test_uuid = "12345678-1234-5678-1234-567812345678"
+        mock_uuid.return_value = uuid.UUID(test_uuid)
+        mock_file_open.side_effect = PermissionError("Permission denied")
+
+        # Execute and Assert
+        with pytest.raises(PermissionError):
+            FileService.save_file(
+                content="Hello World", file_name="test.txt", save_dir="test_dir"
+            )
+
+    @patch("uuid.uuid4")
+    @patch("os.path.getsize")
+    @patch("builtins.open")
+    @patch("os.makedirs")
+    def test_save_file_io_error(
+        self, mock_makedirs, mock_file_open, mock_getsize, mock_uuid
+    ):
+        # Setup
+        test_uuid = "12345678-1234-5678-1234-567812345678"
+        mock_uuid.return_value = uuid.UUID(test_uuid)
+        mock_file_open.side_effect = OSError("IO Error")
+
+        # Execute and Assert
+        with pytest.raises(IOError):
+            FileService.save_file(
+                content="Hello World", file_name="test.txt", save_dir="test_dir"
+            )

llama-index-server/tests/test_llamaindex_server.py

@@ -0,0 +1,106 @@
+import pytest
+from httpx import ASGITransport, AsyncClient
+
+from llama_index.core.agent.workflow import AgentWorkflow
+from llama_index.core.llms import MockLLM
+from llama_index.server import LlamaIndexServer
+
+
+def fetch_weather(city: str) -> str:
+    """Fetch the weather for a given city."""
+    return f"The weather in {city} is sunny."
+
+
+def _agent_workflow() -> AgentWorkflow:
+    # Use MockLLM instead of default OpenAI
+    mock_llm = MockLLM()
+    return AgentWorkflow.from_tools_or_functions(
+        tools_or_functions=[fetch_weather],
+        verbose=True,
+        llm=mock_llm,
+    )
+
+
+@pytest.fixture()
+def server() -> LlamaIndexServer:
+    """Fixture to create a LlamaIndexServer instance."""
+    return LlamaIndexServer(
+        workflow_factory=_agent_workflow,
+        verbose=True,
+        use_default_routers=True,
+        mount_ui=False,
+        env="dev",
+    )
+
+
+@pytest.mark.asyncio()
+async def test_server_has_chat_route(server: LlamaIndexServer) -> None:
+    """Test that the server has the chat API route."""
+    chat_route_exists = any(route.path == "/api/chat" for route in server.routes)
+    assert chat_route_exists, "Chat API route not found in server routes"
+
+
+@pytest.mark.asyncio()
+async def test_server_swagger_docs(server: LlamaIndexServer) -> None:
+    """Test that the server serves Swagger UI docs."""
+    async with AsyncClient(
+        transport=ASGITransport(app=server), base_url="http://test"
+    ) as ac:
+        response = await ac.get("/docs")
+        assert response.status_code == 200
+        assert "text/html" in response.headers["content-type"]
+        assert "Swagger UI" in response.text
+
+
+@pytest.mark.asyncio()
+async def test_ui_is_downloaded(server: LlamaIndexServer) -> None:
+    """
+    Test if the UI is downloaded and mounted correctly.
+    """
+    import os
+    import shutil
+
+    # Clean up any existing static directory first
+    if os.path.exists(".ui"):
+        shutil.rmtree(".ui")
+
+    # Create a new server with UI enabled
+    ui_server = LlamaIndexServer(
+        workflow_factory=_agent_workflow,
+        verbose=True,
+        use_default_routers=True,
+        env="dev",
+        include_ui=True,
+    )
+
+    # Verify that static directory was created with index.html
+    assert os.path.exists("./.ui"), "Static directory was not created"
+    assert os.path.isdir("./.ui"), "Static path is not a directory"
+    assert os.path.exists("./.ui/index.html"), "index.html was not downloaded"
+
+    # Check if the UI is mounted and accessible
+    async with AsyncClient(
+        transport=ASGITransport(app=ui_server), base_url="http://test"
+    ) as ac:
+        response = await ac.get("/")
+        assert response.status_code == 200
+        assert "text/html" in response.headers["content-type"]
+
+    # Clean up after test
+    shutil.rmtree("./.ui")
+
+
+@pytest.mark.asyncio()
+async def test_ui_is_accessible(server: LlamaIndexServer) -> None:
+    """
+    Test if the UI is accessible.
+    """
+    # Manually trigger UI mounting
+    server.mount_ui()
+
+    async with AsyncClient(
+        transport=ASGITransport(app=server), base_url="http://test"
+    ) as ac:
+        response = await ac.get("/")
+        assert response.status_code == 200
+        assert "text/html" in response.headers["content-type"]

llama-index-server/tests/tools/test_document_generator.py

@@ -0,0 +1,89 @@
+import os
+from io import BytesIO
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from llama_index.server.tools.document_generator import (
+    OUTPUT_DIR,
+    DocumentGenerator,
+)
+
+
+class TestDocumentGenerator:
+    @pytest.fixture()
+    def env_setup(self):  # type: ignore
+        os.environ["FILESERVER_URL_PREFIX"] = "http://test-server"
+        yield
+        os.environ.pop("FILESERVER_URL_PREFIX", None)
+
+    def test_validate_file_name(self) -> None:
+        # Valid names
+        assert DocumentGenerator._validate_file_name("valid-name") == "valid-name"
+
+        # Invalid names
+        with pytest.raises(ValueError):
+            DocumentGenerator._validate_file_name("/invalid/path")
+
+    @patch("os.makedirs")
+    @patch("builtins.open")
+    def test_write_to_file(self, mock_open, mock_makedirs):  # type: ignore
+        content = BytesIO(b"test")
+        DocumentGenerator._write_to_file(content, "path/file.txt")
+
+        mock_makedirs.assert_called_once()
+        mock_open.assert_called_once()
+        mock_open.return_value.__enter__.return_value.write.assert_called_once_with(
+            b"test"
+        )
+
+    @patch("markdown.markdown")
+    def test_html_generation(self, mock_markdown):  # type: ignore
+        mock_markdown.return_value = "<h1>Test</h1>"
+
+        # Test HTML content generation
+        assert DocumentGenerator._generate_html_content("# Test") == "<h1>Test</h1>"
+
+        # Test full HTML generation
+        html = DocumentGenerator._generate_html("<h1>Test</h1>")
+        assert "<!DOCTYPE html>" in html
+        assert "<h1>Test</h1>" in html
+
+    @patch("xhtml2pdf.pisa.pisaDocument")
+    def test_pdf_generation(self, mock_pisa):  # type: ignore
+        # Success case
+        mock_pisa.return_value = MagicMock(err=None)
+        assert isinstance(DocumentGenerator._generate_pdf("test"), BytesIO)
+
+        # Error case
+        mock_pisa.return_value = MagicMock(err="Error")
+        with pytest.raises(ValueError):
+            DocumentGenerator._generate_pdf("test")
+
+    @patch.multiple(
+        DocumentGenerator,
+        _generate_html_content=MagicMock(return_value="<h1>Test</h1>"),
+        _generate_html=MagicMock(
+            return_value="<html><body><h1>Test</h1></body></html>"
+        ),
+        _generate_pdf=MagicMock(return_value=BytesIO(b"pdf")),
+        _write_to_file=MagicMock(),
+    )
+    def test_generate_document(self, env_setup):  # type: ignore
+        # HTML generation
+        url = DocumentGenerator.generate_document("# Test", "html", "test-doc")
+        assert url == f"http://test-server/{OUTPUT_DIR}/test-doc.html"
+
+        # PDF generation
+        url = DocumentGenerator.generate_document("# Test", "pdf", "test-doc")
+        assert url == f"http://test-server/{OUTPUT_DIR}/test-doc.pdf"
+
+        # Invalid type
+        with pytest.raises(ValueError):
+            DocumentGenerator.generate_document("# Test", "invalid", "test-doc")
+
+    def test_to_tool(self):  # type: ignore
+        tool = DocumentGenerator().to_tool()
+        # Check the function is correct
+        assert tool.fn == DocumentGenerator.generate_document
+        assert callable(tool.fn)

llama-index-server/tests/tools/test_interpreter.py

@@ -0,0 +1,68 @@
+import os
+from unittest.mock import MagicMock, patch
+
+import pytest
+from e2b_code_interpreter.models import Execution, Logs
+
+from llama_index.server.tools.interpreter import E2BCodeInterpreter
+
+
+class TestE2BCodeInterpreter:
+    @pytest.fixture()
+    def sandbox(self):  # type: ignore
+        """Create a mock Sandbox with no API key requirement."""
+        mock_sandbox = MagicMock()
+        mock_sandbox.files = MagicMock()
+        mock_sandbox.files.write = MagicMock()
+        mock_sandbox.run_code = MagicMock()
+        return mock_sandbox
+
+    @pytest.fixture()
+    def code_interpreter(self, sandbox):  # type: ignore
+        """Create E2BCodeInterpreter that uses the mock Sandbox."""
+        with patch.dict(os.environ, {"E2B_API_KEY": "dummy_key"}):
+            interpreter = E2BCodeInterpreter()
+            interpreter.interpreter = sandbox
+            return interpreter
+
+    def test_interpret_success(self, code_interpreter, sandbox) -> None:  # type: ignore
+        """Test successful code execution."""
+        # Mock execution result
+        mock_execution = Execution()
+        mock_execution.error = None
+        mock_execution.results = []
+        mock_execution.logs = Logs(
+            stdout="stdout", stderr="", display_data="", error=""
+        )
+        sandbox.run_code.return_value = mock_execution
+
+        # Run the code
+        result = code_interpreter.interpret("print('hello')")
+
+        # Verify
+        sandbox.run_code.assert_called_once_with("print('hello')")
+        assert result.is_error is False
+        assert result.logs == mock_execution.logs
+
+    def test_interpret_error(self, code_interpreter, sandbox) -> None:  # type: ignore
+        """Test error in code execution."""
+        # Mock execution result with error
+        mock_execution = Execution()
+        mock_execution.error = "Test error"
+        mock_execution.logs = Logs(
+            stdout="", stderr="error", display_data="", error="Test error"
+        )
+        sandbox.run_code.return_value = mock_execution
+
+        # Run the code
+        result = code_interpreter.interpret("bad code")
+
+        # Verify
+        assert result.is_error is True
+        assert "Error: Test error" in result.error_message
+        sandbox.kill.assert_called_once()
+
+    def test_to_tool(self, code_interpreter) -> None:  # type: ignore
+        """Test tool conversion."""
+        tool = code_interpreter.to_tool()
+        assert tool.fn == code_interpreter.interpret