Merge pull request #61 from GunaPalanivel/feature/19-file-upload-support

Add file upload support for chat sessions (#19)

Author: cnu1812

chatbot-core/api/models/schemas.py

@@ -6,7 +6,30 @@
 """
 
 from enum import Enum
-from pydantic import BaseModel, field_validator
+from typing import List, Optional
+from pydantic import BaseModel, field_validator, model_validator
+
+
+class FileType(str, Enum):
+    """Enum representing supported file types."""
+    TEXT = "text"
+    IMAGE = "image"
+
+
+class FileAttachment(BaseModel):
+    """
+    Represents a processed file attachment.
+
+    Fields:
+        filename (str): Original name of the uploaded file.
+        type (FileType): Type of file - TEXT or IMAGE.
+        content (str): Text content or base64 encoded image data.
+        mime_type (str): MIME type of the file.
+    """
+    filename: str
+    type: FileType
+    content: str
+    mime_type: str
 
 
 class ChatRequest(BaseModel):
@@ -28,12 +51,80 @@ def message_must_not_be_empty(cls, v): # pylint: disable=no-self-argument
             raise ValueError("Message cannot be empty.")
         return v
 
+
+class ChatRequestWithFiles(BaseModel):
+    """
+    Represents a user message with optional file attachments.
+
+    Fields:
+        message (str): The user's input message.
+        files (List[FileAttachment]): Optional list of file attachments.
+
+    Validation:
+        - Rejects when both message is empty and no files are attached.
+    """
+    message: str = ""
+    files: Optional[List[FileAttachment]] = None
+
+    @model_validator(mode="after")
+    def validate_message_or_files(self):
+        """Validates that at least message or files are present."""
+        has_message = bool(self.message and self.message.strip())
+        has_files = bool(self.files and len(self.files) > 0)
+        if not has_message and not has_files:
+            raise ValueError("Either message or files must be provided.")
+        return self
+
 class ChatResponse(BaseModel):
     """
     Represents the chatbot's reply.
     """
     reply: str
 
+
+class ChatResponseWithFiles(BaseModel):
+    """
+    Represents the chatbot's reply with information about processed files.
+
+    Fields:
+        reply (str): The chatbot's text response.
+        processed_files (List[str]): List of filenames that were processed.
+    """
+    reply: str
+    processed_files: Optional[List[str]] = None
+
+
+class FileUploadResponse(BaseModel):
+    """
+    Response model for file upload operations.
+
+    Fields:
+        success (bool): Whether the upload was successful.
+        filename (str): Name of the uploaded file.
+        type (str): Type of file processed ("text" or "image").
+        message (str): Status message.
+    """
+    success: bool
+    filename: str
+    type: str
+    message: str
+
+
+class SupportedExtensionsResponse(BaseModel):
+    """
+    Response model for supported file extensions.
+
+    Fields:
+        text (List[str]): List of supported text file extensions.
+        image (List[str]): List of supported image file extensions.
+        max_text_size_mb (float): Maximum text file size in MB.
+        max_image_size_mb (float): Maximum image file size in MB.
+    """
+    text: List[str]
+    image: List[str]
+    max_text_size_mb: float
+    max_image_size_mb: float
+
 class SessionResponse(BaseModel):
     """
     Response model when a new chat session is created.

chatbot-core/api/routes/chatbot.py

@@ -6,19 +6,27 @@
 the chat service logic.
 """
 
-from fastapi import APIRouter, HTTPException, Response, status
+from typing import List, Optional
+from fastapi import APIRouter, HTTPException, Response, status, UploadFile, File, Form
 from api.models.schemas import (
     ChatRequest,
     ChatResponse,
     SessionResponse,
-    DeleteResponse
+    DeleteResponse,
+    FileAttachment,
+    SupportedExtensionsResponse
 )
 from api.services.chat_service import get_chatbot_reply
 from api.services.memory import (
     init_session,
     delete_session,
     session_exists
 )
+from api.services.file_service import (
+    process_uploaded_file,
+    get_supported_extensions,
+    FileProcessingError
+)
 
 router = APIRouter()
 
@@ -61,6 +69,90 @@ def chatbot_reply(session_id: str, request: ChatRequest):
     return get_chatbot_reply(session_id, request.message)
 
 
+@router.post("/sessions/{session_id}/message/upload", response_model=ChatResponse)
+async def chatbot_reply_with_files(
+    session_id: str,
+    message: str = Form(...),
+    files: Optional[List[UploadFile]] = File(None)
+):
+    """
+    POST endpoint to handle chatbot replies with file uploads.
+
+    Receives a user message with optional file attachments and returns
+    the assistant's reply. Files are processed and their content is
+    included in the context for the LLM.
+
+    Supported file types:
+    - Text files: .txt, .log, .md, .json, .xml, .yaml, .yml, code files
+    - Image files: .png, .jpg, .jpeg, .gif, .webp, .bmp
+
+    Args:
+        session_id (str): The ID of the session from the URL path.
+        message (str): The user's message (form field).
+        files (List[UploadFile]): Optional list of uploaded files.
+
+    Returns:
+        ChatResponse: The chatbot's generated reply.
+
+    Raises:
+        HTTPException: 404 if session not found, 400 if file processing fails,
+                      422 if message is empty and no files provided.
+    """
+    if not session_exists(session_id):
+        raise HTTPException(status_code=404, detail="Session not found.")
+
+    # Validate that at least message or files are provided
+    has_message = message and message.strip()
+    has_files = files and len(files) > 0
+
+    if not has_message and not has_files:
+        raise HTTPException(
+            status_code=422,
+            detail="Either message or files must be provided."
+        )
+
+    # Process uploaded files
+    processed_files: List[FileAttachment] = []
+
+    if files:
+        for upload_file in files:
+            try:
+                content = await upload_file.read()
+                processed = process_uploaded_file(
+                    content, upload_file.filename or "unknown"
+                )
+                processed_files.append(FileAttachment(**processed))
+            except FileProcessingError as e:
+                raise HTTPException(status_code=400, detail=str(e)) from e
+            except Exception as e:
+                raise HTTPException(
+                    status_code=500,
+                    detail=f"Failed to process file: {type(e).__name__}"
+                ) from e
+            finally:
+                await upload_file.close()
+
+    # Use default message if only files provided
+    final_message = message.strip() if has_message else "Please analyze the attached file(s)."
+
+    return get_chatbot_reply(
+        session_id, final_message, processed_files if processed_files else None
+    )
+
+
+@router.get("/files/supported-extensions", response_model=SupportedExtensionsResponse)
+def get_supported_file_extensions():
+    """
+    GET endpoint to retrieve supported file extensions for upload.
+
+    Returns:
+        SupportedExtensionsResponse: Lists of supported text and image extensions,
+                                     along with size limits.
+    """
+    extensions = get_supported_extensions()
+    return SupportedExtensionsResponse(**extensions)
+
+
 @router.delete("/sessions/{session_id}", response_model=DeleteResponse)
 def delete_chat(session_id: str):
     """

chatbot-core/api/services/chat_service.py

@@ -15,8 +15,9 @@
     RETRIEVER_AGENT_PROMPT,
     CONTEXT_RELEVANCE_PROMPT
 )
-from api.models.schemas import ChatResponse, QueryType, try_str_to_query_type
+from api.models.schemas import ChatResponse, QueryType, try_str_to_query_type, FileAttachment
 from api.services.memory import get_session
+from api.services.file_service import format_file_context
 from api.models.embedding_model import EMBEDDING_MODEL
 from api.tools.tools import TOOL_REGISTRY
 from api.tools.utils import get_default_tools_call, validate_tool_calls, make_placeholder_replacer
@@ -28,34 +29,57 @@
 retrieval_config = CONFIG["retrieval"]
 CODE_BLOCK_PLACEHOLDER_PATTERN = r"\[\[(?:CODE_BLOCK|CODE_SNIPPET)_(\d+)\]\]"
 
-def get_chatbot_reply(session_id: str, user_input: str) -> ChatResponse:
+def get_chatbot_reply(
+    session_id: str,
+    user_input: str,
+    files: Optional[List[FileAttachment]] = None
+) -> ChatResponse:
     """
     Main chatbot entry point. Retrieves context, constructs a prompt with memory,
     and generates an LLM response. Also updates the memory with the latest exchange.
 
     Args:
         session_id (str): The unique ID for the chat session.
         user_input (str): The latest user message.
+        files (Optional[List[FileAttachment]]): Optional list of file attachments.
 
     Returns:
         ChatResponse: The generated assistant response.
     """
     logger.info("New message from session '%s'", session_id)
     logger.info("Handling the user query: %s", user_input)
 
+    if files:
+        logger.info("Processing %d uploaded file(s)", len(files))
+
     memory = get_session(session_id)
     if memory is None:
         raise RuntimeError(f"Session '{session_id}' not found in the memory store.")
 
     context = retrieve_context(user_input)
     logger.info("Context retrieved: %s", context)
 
+    # Add file context if files are provided
+    file_context = ""
+    if files:
+        file_dicts = [file.model_dump() for file in files]
+        file_context = format_file_context(file_dicts)
+        if file_context:
+            logger.info("File context added: %d characters", len(file_context))
+            context = f"{context}\n\n[User Uploaded Files]\n{file_context}"
+
     prompt = build_prompt(user_input, context, memory)
 
     logger.info("Generating answer with prompt: %s", prompt)
     reply = generate_answer(prompt)
 
-    memory.chat_memory.add_user_message(user_input)
+    # Include file info in memory message
+    user_message = user_input
+    if files:
+        file_names = [f.filename for f in files]
+        user_message = f"{user_input}\n[Attached files: {', '.join(file_names)}]"
+
+    memory.chat_memory.add_user_message(user_message)
     memory.chat_memory.add_ai_message(reply)
 
     return ChatResponse(reply=reply)