feat: Support citation for agentic template

.changeset/few-news-marry.md

@@ -0,0 +1,5 @@
+---
+"@create-llama/llama-index-server": patch
+---
+
+Show agent widget in UI when making a tool call

.changeset/rich-nights-hug.md

@@ -0,0 +1,5 @@
+---
+"@create-llama/llama-index-server": patch
+---
+
+Support citation for query engine tool

.changeset/small-insects-hug.md

@@ -0,0 +1,5 @@
+---
+"create-llama": patch
+---
+
+Support citation for agentic template (Python)

.changeset/stale-things-beg.md

@@ -0,0 +1,5 @@
+---
+"@llamaindex/server": patch
+---
+
+Bump version: [email protected]

.github/workflows/e2e.yml

@@ -64,6 +64,15 @@ jobs:
         run: pnpm run pack-install
         working-directory: packages/create-llama
 
+      - name: Build and store server package
+        run: |
+          pnpm run build
+          wheel_file=$(ls dist/*.whl | head -n 1)
+          mkdir -p "${{ runner.temp }}"
+          cp "$wheel_file" "${{ runner.temp }}/"
+          echo "SERVER_PACKAGE_PATH=${{ runner.temp }}/$(basename "$wheel_file")" >> $GITHUB_ENV
+        working-directory: python/llama-index-server
+
       - name: Run Playwright tests for Python
         run: pnpm run e2e:python
         env:
@@ -74,6 +83,7 @@ jobs:
           TEMPLATE_TYPE: ${{ matrix.template-types }}
           PYTHONIOENCODING: utf-8
           PYTHONLEGACYWINDOWSSTDIO: utf-8
+          SERVER_PACKAGE_PATH: ${{ env.SERVER_PACKAGE_PATH }}
         working-directory: packages/create-llama
 
       - uses: actions/upload-artifact@v4

packages/create-llama/helpers/python.ts

@@ -5,6 +5,7 @@ import { parse, stringify } from "smol-toml";
 import terminalLink from "terminal-link";
 import { isUvAvailable, tryUvSync } from "./uv";
 
+import { isCI } from "ci-info";
 import { assetRelocator, copy } from "./copy";
 import { templatesDir } from "./dir";
 import { Tool } from "./tools";
@@ -278,6 +279,19 @@ const getAdditionalDependencies = (
     }
   }
 
+  // If app template is llama-index-server and CI and SERVER_PACKAGE_PATH is set,
+  // add @llamaindex/server to dependencies
+  if (
+    templateType === "llamaindexserver" &&
+    isCI &&
+    process.env.SERVER_PACKAGE_PATH
+  ) {
+    dependencies.push({
+      name: "llama-index-server",
+      version: `@file://${process.env.SERVER_PACKAGE_PATH}`,
+    });
+  }
+
   return dependencies;
 };
 

packages/create-llama/templates/components/use-cases/python/agentic_rag/workflow.py

@@ -3,9 +3,12 @@
 from app.index import get_index
 from llama_index.core.agent.workflow import AgentWorkflow
 from llama_index.core.settings import Settings
-from llama_index.llms.openai import OpenAI
 from llama_index.server.api.models import ChatRequest
 from llama_index.server.tools.index import get_query_engine_tool
+from llama_index.server.tools.index.citation import (
+    CITATION_SYSTEM_PROMPT,
+    enable_citation,
+)
 
 
 def create_workflow(chat_request: Optional[ChatRequest] = None) -> AgentWorkflow:
@@ -14,9 +17,16 @@ def create_workflow(chat_request: Optional[ChatRequest] = None) -> AgentWorkflow
         raise RuntimeError(
             "Index not found! Please run `uv run generate` to index the data first."
         )
-    query_tool = get_query_engine_tool(index=index)
+    # Create a query tool with citations enabled
+    query_tool = enable_citation(get_query_engine_tool(index=index))
+
+    # Define the system prompt for the agent
+    # Append the citation system prompt to the system prompt
+    system_prompt = """You are a helpful assistant"""
+    system_prompt += CITATION_SYSTEM_PROMPT
+
     return AgentWorkflow.from_tools_or_functions(
         tools_or_functions=[query_tool],
-        llm=Settings.llm or OpenAI(model="gpt-4o-mini"),
-        system_prompt="You are a helpful assistant.",
+        llm=Settings.llm,
+        system_prompt=system_prompt,
     )

packages/create-llama/templates/types/llamaindexserver/fastapi/app/settings.py

@@ -8,5 +8,5 @@
 def init_settings():
     if os.getenv("OPENAI_API_KEY") is None:
         raise RuntimeError("OPENAI_API_KEY is missing in environment variables")
-    Settings.llm = OpenAI(model="gpt-4o-mini")
-    Settings.embed_model = OpenAIEmbedding(model="text-embedding-3-small")
+    Settings.llm = OpenAI(model="gpt-4.1")
+    Settings.embed_model = OpenAIEmbedding(model="text-embedding-3-large")

packages/create-llama/templates/types/llamaindexserver/fastapi/pyproject.toml

@@ -46,6 +46,9 @@ disable_error_code = [ "return-value", "assignment" ]
 module = "app.*"
 ignore_missing_imports = false
 
+[tool.hatch.metadata]
+allow-direct-references = true
+
 [build-system]
 requires = [ "hatchling>=1.24" ]
-build-backend = "hatchling.build"
+build-backend = "hatchling.build"

packages/server/package.json

@@ -59,7 +59,7 @@
     "@babel/traverse": "^7.27.0",
     "@babel/types": "^7.27.0",
     "@hookform/resolvers": "^5.0.1",
-    "@llamaindex/chat-ui": "0.4.5",
+    "@llamaindex/chat-ui": "0.4.6",
     "@radix-ui/react-accordion": "^1.2.3",
     "@radix-ui/react-alert-dialog": "^1.1.7",
     "@radix-ui/react-aspect-ratio": "^1.1.3",

python/llama-index-server/examples/llamacloud/main.py

@@ -0,0 +1,113 @@
+import os
+from typing import List, Optional
+
+from fastapi import FastAPI
+from llama_index.core.agent.workflow import AgentWorkflow
+from llama_index.core.query_engine.retriever_query_engine import RetrieverQueryEngine
+from llama_index.core.settings import Settings
+from llama_index.core.tools import QueryEngineTool, ToolMetadata
+from llama_index.llms.openai import OpenAI
+from llama_index.server import LlamaIndexServer, UIConfig
+from llama_index.server.api.models import ChatRequest
+from llama_index.server.services.llamacloud import LlamaCloudIndex, get_index
+from llama_index.server.tools.index.citation import (
+    CITATION_SYSTEM_PROMPT,
+    enable_citation,
+)
+
+# Please set the following environment variables to use LlamaCloud
+if os.getenv("LLAMA_CLOUD_API_KEY") is None:
+    raise ValueError("LLAMA_CLOUD_API_KEY is not set")
+if os.getenv("LLAMA_CLOUD_PROJECT_NAME") is None:
+    raise ValueError("LLAMA_CLOUD_PROJECT_NAME is not set")
+if os.getenv("LLAMA_CLOUD_INDEX_NAME") is None:
+    raise ValueError("LLAMA_CLOUD_INDEX_NAME is not set")
+
+Settings.llm = OpenAI(model="gpt-4.1")
+
+
+def get_tools(index: LlamaCloudIndex) -> List[QueryEngineTool]:
+    """
+    Get the tools for the given index.
+    """
+
+    chunk_retriever = index.as_retriever(
+        retrieval_mode="chunks",
+        rerank_top_n=15,
+        dense_similarity_top_k=1,
+    )
+    doc_retriever = index.as_retriever(
+        retrieval_mode="files_via_content",
+        files_top_k=1,
+    )
+
+    # You can either create query engine with CitationSynthesizer and NodeCitationProcessor
+    # or use the enable_citation function to enable citation for the query engine.
+    chunk_engine = RetrieverQueryEngine.from_args(
+        retriever=chunk_retriever,
+        llm=Settings.llm,
+    )
+    doc_engine = RetrieverQueryEngine.from_args(
+        retriever=doc_retriever,
+        llm=Settings.llm,
+    )
+
+    chunk_tool = QueryEngineTool(
+        query_engine=chunk_engine,
+        metadata=ToolMetadata(
+            name="chunk_query_engine",
+            description=(
+                "Get answer from specific chunk of a given document. Best used for lower-level questions that require specific information from a given document."
+                "Do NOT use if the answer can be found in the entire document. Use the file_query_engine instead for that purpose"
+            ),
+        ),
+    )
+    doc_tool = QueryEngineTool(
+        query_engine=doc_engine,
+        metadata=ToolMetadata(
+            name="file_query_engine",
+            description=(
+                "Get answer from entire document as context.  Best used for higher-level summarization questions."
+                "Do NOT use if the answer can be found in a specific chunk of a given document. Use the chunk_query_engine instead for that purpose"
+            ),
+        ),
+    )
+
+    return [enable_citation(chunk_tool), enable_citation(doc_tool)]
+
+
+def create_workflow(chat_request: Optional[ChatRequest] = None) -> AgentWorkflow:
+    index = get_index(chat_request=chat_request)
+    if index is None:
+        raise RuntimeError("Index not found!")
+
+    # Append the citation system prompt to the system prompt
+    system_prompt = """
+    You are a helpful assistant that has access to a knowledge base.
+    """
+    system_prompt += CITATION_SYSTEM_PROMPT
+    return AgentWorkflow.from_tools_or_functions(
+        tools_or_functions=get_tools(index),
+        system_prompt=system_prompt,
+    )
+
+
+def create_app() -> FastAPI:
+    app = LlamaIndexServer(
+        workflow_factory=create_workflow,
+        env="dev",
+        suggest_next_questions=False,
+        ui_config=UIConfig(
+            llamacloud_index_selector=True,  # to select different indexes in the UI
+        ),
+    )
+    return app
+
+
+app = create_app()
+
+
+if __name__ == "__main__":
+    import uvicorn
+
+    uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True)

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

@@ -1,3 +1,4 @@
+from llama_index.server.api.callbacks.agent_call_tool import AgentCallTool
 from llama_index.server.api.callbacks.base import EventCallback
 from llama_index.server.api.callbacks.llamacloud import LlamaCloudFileDownload
 from llama_index.server.api.callbacks.source_nodes import SourceNodesFromToolCall
@@ -10,4 +11,5 @@
     "SourceNodesFromToolCall",
     "SuggestNextQuestions",
     "LlamaCloudFileDownload",
+    "AgentCallTool",
 ]

python/llama-index-server/llama_index/server/api/callbacks/agent_call_tool.py

@@ -0,0 +1,26 @@
+import logging
+from typing import Any
+
+from llama_index.core.agent.workflow.workflow_events import ToolCall, ToolCallResult
+from llama_index.server.api.callbacks.base import EventCallback
+from llama_index.server.api.models import AgentRunEvent
+
+logger = logging.getLogger("uvicorn")
+
+
+class AgentCallTool(EventCallback):
+    """
+    Adapter for convert tool call events to agent run events.
+    """
+
+    async def run(self, event: Any) -> Any:
+        if isinstance(event, ToolCall) and not isinstance(event, ToolCallResult):
+            return AgentRunEvent(
+                name="Agent",
+                msg=f"Calling tool: {event.tool_name} with: {event.tool_kwargs}",
+            )
+        return event
+
+    @classmethod
+    def from_default(cls, *args: Any, **kwargs: Any) -> "AgentCallTool":
+        return cls()

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

@@ -1,31 +1,51 @@
-from typing import Any
+import logging
+from typing import Any, List, Optional
 
 from llama_index.core.agent.workflow.workflow_events import ToolCallResult
+from llama_index.core.schema import NodeWithScore
 from llama_index.server.api.callbacks.base import EventCallback
 from llama_index.server.api.models import SourceNodesEvent
 
+logger = logging.getLogger(__name__)
+
 
 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)
+    def __init__(self, tool_name: Optional[str] = None):
+        # backward compatibility
+        if tool_name is not None:
+            logger.warning(
+                "tool_name has been deprecated. It's now detected by the tool output."
+            )
+
+    def _get_source_nodes(self, event: ToolCallResult) -> Optional[List[NodeWithScore]]:
+        # If result is not error
+        if event.tool_output.is_error:
+            return None
+        # If result is not error, check if source nodes are in the tool output
+        raw_output = event.tool_output.raw_output
+        if hasattr(raw_output, "source_nodes"):
+            source_nodes = raw_output.source_nodes
+            # Verify if source_nodes is List[NodeWithScore]
+            if isinstance(source_nodes, list) and all(
+                isinstance(node, NodeWithScore) for node in source_nodes
+            ):
+                return source_nodes
+            else:
+                return None
+        else:
+            return None
 
     async def run(self, event: Any) -> Any:
+        events = [event]
         if isinstance(event, ToolCallResult):
-            if event.tool_name == self.query_tool_name:
-                return event, self.transform_tool_call_result(event)
-        return event
+            source_nodes = self._get_source_nodes(event)
+            if source_nodes is not None:
+                events.append(SourceNodesEvent(nodes=source_nodes))
+        return events
 
     @classmethod
     def from_default(cls, *args: Any, **kwargs: Any) -> "SourceNodesFromToolCall":