attach file to user message

Author: leehuwuj

packages/server/next/app/api/files/helpers.ts

@@ -8,7 +8,9 @@ export async function storeFile(name: string, fileBuffer: Buffer) {
   const parts = name.split(".");
   const fileName = parts[0];
   const fileExt = parts[1];
-
+  if (!fileName) {
+    throw new Error("File name is required");
+  }
   if (!fileExt) {
     throw new Error("File extension is required");
   }

python/llama-index-server/examples/private_file/agent-workflow.py

@@ -18,7 +18,7 @@ def create_file_tool(chat_request: ChatRequest) -> Optional[FunctionTool]:
     Create a tool to read file if the user uploads a file.
     """
     file_ids = []
-    for file in get_file_attachments(chat_request):
+    for file in get_file_attachments(chat_request.messages):
         file_ids.append(file.id)
     if len(file_ids) == 0:
         return None
@@ -29,7 +29,7 @@ def create_file_tool(chat_request: ChatRequest) -> Optional[FunctionTool]:
     )
 
     def read_file(file_id: str) -> str:
-        file_path = FileService.get_private_file_path(file_id)
+        file_path = FileService.get_file_path(file_id)
         try:
             with open(file_path, "r") as file:
                 return file.read()
@@ -57,7 +57,7 @@ def create_app() -> FastAPI:
         workflow_factory=create_workflow,
         suggest_next_questions=False,
         ui_config=UIConfig(
-            file_upload_enabled=True,
+            enable_file_upload=True,
             component_dir="components",
         ),
     )

python/llama-index-server/examples/private_file/custom-workflow.py

@@ -2,21 +2,19 @@
 
 from fastapi import FastAPI
 from llama_index.core.agent.workflow.workflow_events import AgentStream
-from llama_index.core.llms import LLM
+from llama_index.core.llms import LLM, ChatMessage, DocumentBlock
 from llama_index.core.prompts import PromptTemplate
 from llama_index.core.workflow import (
     Context,
     Event,
     StartEvent,
     StopEvent,
     Workflow,
+    WorkflowRuntimeError,
     step,
 )
 from llama_index.llms.openai import OpenAI
 from llama_index.server import LlamaIndexServer, UIConfig
-from llama_index.server.models import ChatRequest
-from llama_index.server.services.file import FileService
-from llama_index.server.utils.chat_attachments import get_file_attachments
 
 
 class FileHelpEvent(Event):
@@ -25,7 +23,7 @@ class FileHelpEvent(Event):
     """
 
     file_content: str
-    user_msg: str
+    user_request: str
 
 
 class FileHelpWorkflow(Workflow):
@@ -38,53 +36,25 @@ class FileHelpWorkflow(Workflow):
     def __init__(
         self,
         llm: LLM,
-        chat_request: ChatRequest,  # Initial the workflow with the chat request
         **kwargs: Any,
     ):
         super().__init__(**kwargs)
         self.llm = llm
-        # Get the uploaded files from the chat request and stores them in the workflow instance for accessing later
-        self.uploaded_files = get_file_attachments(chat_request)
-        if len(self.uploaded_files) == 0:
-            raise ValueError("No uploaded files found. Please upload a file to start")
 
     @step
     async def read_files(self, ctx: Context, ev: StartEvent) -> FileHelpEvent:
-        user_msg = ev.user_msg
+        user_msg: ChatMessage = ev.user_msg
+        # All the uploaded files are included in the user_msg.blocks as DocumentBlock
+        files = [block for block in user_msg.blocks if isinstance(block, DocumentBlock)]
+        if len(files) != 1:
+            raise WorkflowRuntimeError("Please upload only one file")
 
-        # 1. Access through workflow instance as is
-        # last_file = self.uploaded_files[-1]
-
-
-        # 2. Access through user_msg (if it's a ChatMessage)
-        # llama_index support ChatMessage with DocumentBlock which mostly the same as our FileServer.
-        # (but I guess we'll get back to dealing with other problems
-        # that we need to pass other data to the workflow later)
-        # e.g:
-        # files = [
-        #     ServerFile.from_document_block(block)
-        #     for block in user_msg.blocks
-        #     if isinstance(block, DocumentBlock)
-        # ]
-        #
-        # or they can just use files: List[DocumentBlock] as is.
-
-        
-        # 3. Introduce server start event with additional fields
-        # e.g:
-        # class ChatStartEvent(StartEvent):
-        #     user_msg: Union[str, ChatMessage]
-        #     chat_history: list[ChatMessage]
-        #     attachments: list[ServerFile]
-        # Then the user can clearly know what do they have with the StartEvent
-
-        file_path = FileService.get_private_file_path(last_file.id)
-        with open(file_path, "r", encoding="utf-8") as f:
-            file_content = f.read()
+        # Simply call resolve_document() to get the file content
+        file_content = files[0].resolve_document().read().decode("utf-8")
 
         return FileHelpEvent(
             file_content=file_content,
-            user_msg=user_msg,
+            user_request=ev.user_msg.content,
         )
 
     @step
@@ -100,7 +70,7 @@ async def help_user(self, ctx: Context, ev: FileHelpEvent) -> StopEvent:
         {file_content}
         """)
         prompt = default_prompt.format(
-            user_msg=ev.user_msg,
+            user_msg=ev.user_request,
             file_content=ev.file_content,
         )
         stream = await self.llm.astream_complete(prompt)
@@ -120,10 +90,9 @@ async def help_user(self, ctx: Context, ev: FileHelpEvent) -> StopEvent:
         )
 
 
-def create_workflow(chat_request: ChatRequest) -> Workflow:
+def create_workflow() -> Workflow:
     return FileHelpWorkflow(
         llm=OpenAI(model="gpt-4.1-mini"),
-        chat_request=chat_request,
     )
 
 
@@ -132,7 +101,7 @@ def create_app() -> FastAPI:
         workflow_factory=create_workflow,
         suggest_next_questions=False,
         ui_config=UIConfig(
-            file_upload_enabled=True,
+            enable_file_upload=True,
             component_dir="components",
         ),
     )

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

@@ -6,7 +6,6 @@
 
 from fastapi import APIRouter, BackgroundTasks, HTTPException
 from fastapi.responses import StreamingResponse
-
 from llama_index.core.agent.workflow.workflow_events import (
     AgentInput,
     AgentSetup,
@@ -24,6 +23,7 @@
     SuggestNextQuestions,
 )
 from llama_index.server.api.callbacks.stream_handler import StreamHandler
+from llama_index.server.api.utils.chat_request import prepare_user_message
 from llama_index.server.api.utils.vercel_stream import VercelStreamResponse
 from llama_index.server.models.chat import ChatFile, ChatRequest
 from llama_index.server.models.hitl import HumanInputEvent
@@ -46,7 +46,7 @@ async def chat(
     ) -> StreamingResponse:
         try:
             last_message = request.messages[-1]
-            user_message = last_message.to_llamaindex_message()
+            user_message = prepare_user_message(request)
             chat_history = [
                 message.to_llamaindex_message() for message in request.messages[:-1]
             ]
@@ -68,7 +68,7 @@ async def chat(
                 workflow_handler = workflow.run(ctx=ctx)
             else:
                 workflow_handler = workflow.run(
-                    user_msg=user_message.content,
+                    user_msg=user_message,
                     chat_history=chat_history,
                 )
 

python/llama-index-server/llama_index/server/api/utils/chat_request.py

@@ -1,7 +1,11 @@
 from typing import List, Optional
 
+from llama_index.core.llms import DocumentBlock
+from llama_index.core.types import ChatMessage, MessageRole
 from llama_index.server.models.artifacts import Artifact
 from llama_index.server.models.chat import ChatRequest
+from llama_index.server.services.file import FileService
+from llama_index.server.utils.chat_attachments import get_file_attachments
 
 
 def get_artifacts(chat_request: ChatRequest) -> List[Artifact]:
@@ -22,3 +26,25 @@ def get_artifacts(chat_request: ChatRequest) -> List[Artifact]:
 def get_last_artifact(chat_request: ChatRequest) -> Optional[Artifact]:
     artifacts = get_artifacts(chat_request)
     return artifacts[-1] if len(artifacts) > 0 else None
+
+
+def prepare_user_message(chat_request: ChatRequest) -> ChatMessage:
+    """
+    Prepare the user message from the chat request.
+    """
+    last_message: ChatMessage = chat_request.messages[-1].to_llamaindex_message()
+    if last_message.role != MessageRole.USER:
+        raise ValueError("Last message must be from user")
+
+    # Add attached files to the user message
+    attachment_files = get_file_attachments(chat_request.messages)
+    last_message.blocks += [
+        DocumentBlock(
+            path=file.path or FileService.get_file_path(file.id),
+            url=file.url,
+            document_mimetype=file.type,
+        )
+        for file in attachment_files
+    ]
+
+    return last_message

python/llama-index-server/llama_index/server/models/chat.py

@@ -1,9 +1,8 @@
 import re
 from typing import Any, List, Literal, Optional, Union
 
-from pydantic import BaseModel, Field, field_validator
-
 from llama_index.core.types import ChatMessage, MessageRole
+from pydantic import BaseModel, Field, field_validator
 
 
 class ServerFile(BaseModel):
@@ -74,6 +73,10 @@ class ChatAPIMessage(BaseModel):
     annotations: Optional[List[Union[FileAnnotation, Any]]] = None
 
     def to_llamaindex_message(self) -> ChatMessage:
+        """
+        Simply convert text content of API message to llama_index's ChatMessage.
+        Annotations are not included.
+        """
         return ChatMessage(role=self.role, content=self.content)
 
     @property

python/llama-index-server/llama_index/server/utils/chat_attachments.py

@@ -1,23 +1,26 @@
 from typing import List
 
-from llama_index.server.models.chat import ChatRequest, FileAnnotation, ServerFile
+from llama_index.core.types import MessageRole
+from llama_index.server.models.chat import ChatAPIMessage, FileAnnotation, ServerFile
 
 
-def get_file_attachments(chat_request: ChatRequest) -> List[ServerFile]:
+def get_file_attachments(messages: List[ChatAPIMessage]) -> List[ServerFile]:
     """
-    Extract all file attachments from the chat request.
+    Extract all file attachments from user messages.
 
     Args:
-        chat_request (ChatRequest): The chat request.
+        messages (List[ChatAPIMessage]): The list of messages.
 
     Returns:
-        List[PrivateFile]: The list of private files.
+        List[ServerFile]: The list of private files.
     """
-    message_annotations = [
-        message.annotations for message in chat_request.messages if message.annotations
+    user_message_annotations = [
+        message.annotations
+        for message in messages
+        if message.annotations and message.role == MessageRole.USER
     ]
     files: List[ServerFile] = []
-    for annotation in message_annotations:
+    for annotation in user_message_annotations:
         if isinstance(annotation, list):
             for item in annotation:
                 if isinstance(item, FileAnnotation):

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

@@ -90,7 +90,6 @@ async def test_chat_router(
 
         # 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
 
@@ -153,6 +152,5 @@ def workflow_factory(verbose: bool = False) -> MagicMock:
 
         # 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