backend: upload + chat with compass

src/backend/chat/custom/custom.py

@@ -257,8 +257,9 @@ async def call_tools(
                 model_deployment=deployment_model,
                 user_id=ctx.get_user_id(),
                 trace_id=ctx.get_trace_id(),
-                agent_id=ctx.get_agent_id(),
                 agent_tool_metadata=ctx.get_agent_tool_metadata(),
+                agent_id=ctx.get_agent_id(),
+                conversation_id=ctx.get_conversation_id(),
             )
 
             # If the tool returns a list of outputs, append each output to the tool_results list
@@ -302,7 +303,7 @@ def add_files_to_chat_history(
             num_words = min(25, word_count)
             preview = " ".join(file.file_content.split()[:num_words])
 
-            files_message += f"Filename: {file.file_name}\nWord Count: {word_count} Preview: {preview}\n\n"
+            files_message += f"Filename: {file.file_name}\nFile ID: {file.id}\nWord Count: {word_count} Preview: {preview}\n\n"
 
         chat_history.append(ChatMessage(message=files_message, role=ChatRole.SYSTEM))
         return chat_history

src/backend/config/settings.py

@@ -80,6 +80,12 @@ class FeatureFlags(BaseSettings, BaseModel):
             "USE_COMMUNITY_FEATURES", "use_community_features"
         ),
     )
+    use_compass_file_storage: Optional[bool] = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            "USE_COMPASS_FILE_STORAGE", "use_compass_file_storage"
+        ),
+    )
 
 
 class PythonToolSettings(BaseSettings, BaseModel):

src/backend/config/tools.py

@@ -66,9 +66,9 @@ class ToolName(StrEnum):
                 "type": "str",
                 "required": True,
             },
-            "filenames": {
-                "description": "A list of one or more uploaded filename strings to search over",
-                "type": "list",
+            "files": {
+                "description": "A list of files represented as tuples of (filename, file ID) to search over",
+                "type": "list[tuple[str, str]]",
                 "required": True,
             },
         },
@@ -82,9 +82,9 @@ class ToolName(StrEnum):
         display_name="Read Document",
         implementation=ReadFileTool,
         parameter_definitions={
-            "filename": {
-                "description": "The name of the attached file to read.",
-                "type": "str",
+            "file": {
+                "description": "A file represented as a tuple (filename, file ID) to read over",
+                "type": "tuple[str, str]",
                 "required": True,
             }
         },

src/backend/routers/agent.py

@@ -1,6 +1,11 @@
-from fastapi import APIRouter, Depends, HTTPException
+from fastapi import APIRouter, Depends
+from fastapi import File as RequestFile
+from fastapi import HTTPException
+from fastapi import UploadFile as FastAPIUploadFile
 
 from backend.config.routers import RouterName
+from backend.config.settings import Settings
+from backend.config.tools import ToolName
 from backend.crud import agent as agent_crud
 from backend.crud import agent_tool_metadata as agent_tool_metadata_crud
 from backend.database_models.agent import Agent as AgentModel
@@ -21,6 +26,7 @@
     UpdateAgentToolMetadataRequest,
 )
 from backend.schemas.context import Context
+from backend.schemas.file import DeleteFileResponse, UploadFileResponse
 from backend.schemas.metrics import (
     DEFAULT_METRICS_AGENT,
     GenericResponseMessage,
@@ -33,6 +39,11 @@
     validate_agent_tool_metadata_exists,
 )
 from backend.services.context import get_context
+from backend.services.file import (
+    consolidate_agent_files_in_compass,
+    get_file_service,
+    validate_batch_file_size,
+)
 from backend.services.request_validators import (
     validate_create_agent_request,
     validate_update_agent_request,
@@ -98,6 +109,30 @@ async def create_agent(
                 await update_or_create_tool_metadata(
                     created_agent, tool_metadata, session, ctx
                 )
+
+        # Consolidate agent files into one index in compass
+        if (
+            Settings().feature_flags.use_compass_file_storage
+            and created_agent.tools_metadata
+        ):
+            artifacts = next(
+                (
+                    tool_metadata.artifacts
+                    for tool_metadata in created_agent.tools_metadata
+                    if tool_metadata.tool_name == ToolName.Read_File
+                    or tool_metadata.tool_name == ToolName.Search_File
+                ),
+                [],
+            )
+            file_ids = list(
+                set(
+                    artifact.get("id")
+                    for artifact in artifacts
+                    if artifact.get("type") == "local_file"
+                )
+            )
+            if len(file_ids) > 0:
+                await consolidate_agent_files_in_compass(file_ids, created_agent.id)
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
@@ -476,6 +511,62 @@ async def delete_agent_tool_metadata(
     return DeleteAgentToolMetadata()
 
 
+@router.post("/batch_upload_file", response_model=list[UploadFileResponse])
+async def batch_upload_file(
+    session: DBSessionDep,
+    files: list[FastAPIUploadFile] = RequestFile(...),
+    ctx: Context = Depends(get_context),
+) -> UploadFileResponse:
+    user_id = ctx.get_user_id()
+    validate_batch_file_size(session, user_id, files)
+
+    uploaded_files = []
+    try:
+        uploaded_files = await get_file_service().create_agent_files(
+            session,
+            files,
+            user_id,
+            ctx,
+        )
+    except Exception as e:
+        raise HTTPException(
+            status_code=500, detail=f"Error while uploading agent file(s): {e}."
+        )
+
+    return uploaded_files
+
+
+@router.delete("/{agent_id}/files/{file_id}")
+async def delete_file(
+    agent_id: str,
+    file_id: str,
+    session: DBSessionDep,
+    ctx: Context = Depends(get_context),
+) -> DeleteFileResponse:
+    """
+    Delete a file by ID.
+
+    Args:
+        agent_id (str): Agent ID.
+        file_id (str): File ID.
+        session (DBSessionDep): Database session.
+
+    Returns:
+        DeleteFile: Empty response.
+
+    Raises:
+        HTTPException: If the agent with the given ID is not found.
+    """
+    user_id = ctx.get_user_id()
+    _ = validate_agent_exists(session, agent_id)
+    validate_file(session, file_id, user_id, agent_id)
+
+    # Delete the File DB object
+    get_file_service().delete_agent_file_by_id(session, agent_id, file_id, user_id)
+
+    return DeleteFileResponse()
+
+
 # Default Agent Router
 default_agent_router = APIRouter(
     prefix="/v1/default_agent",

src/backend/routers/chat.py

@@ -78,7 +78,6 @@ async def chat_stream(
     (
         session,
         chat_request,
-        file_paths,
         response_message,
         should_store,
         managed_tools,
@@ -92,7 +91,6 @@ async def chat_stream(
             CustomChat().chat(
                 chat_request,
                 stream=True,
-                file_paths=file_paths,
                 managed_tools=managed_tools,
                 session=session,
                 ctx=ctx,
@@ -152,7 +150,6 @@ async def chat(
     (
         session,
         chat_request,
-        file_paths,
         response_message,
         should_store,
         managed_tools,
@@ -165,7 +162,6 @@ async def chat(
         CustomChat().chat(
             chat_request,
             stream=False,
-            file_paths=file_paths,
             managed_tools=managed_tools,
             ctx=ctx,
         ),

src/backend/routers/conversation.py

@@ -6,6 +6,7 @@
 from backend.chat.custom.custom import CustomChat
 from backend.chat.custom.utils import get_deployment
 from backend.config.routers import RouterName
+from backend.config.settings import Settings
 from backend.crud import agent as agent_crud
 from backend.crud import conversation as conversation_crud
 from backend.database_models import Conversation as ConversationModel
@@ -26,6 +27,7 @@
     UploadFileResponse,
 )
 from backend.schemas.metrics import DEFAULT_METRICS_AGENT, agent_to_metrics_agent
+from backend.services.compass import Compass
 from backend.services.context import get_context
 from backend.services.conversation import (
     DEFAULT_TITLE,
@@ -230,12 +232,11 @@ async def delete_conversation(
         HTTPException: If the conversation with the given ID is not found.
     """
     user_id = ctx.get_user_id()
-    _ = validate_conversation(session, conversation_id, user_id)
-    conversation = conversation_crud.get_conversation(session, conversation_id, user_id)
-
-    if conversation.file_ids:
-        get_file_service().bulk_delete_files(session, conversation.file_ids, user_id)
+    conversation = validate_conversation(session, conversation_id, user_id)
 
+    get_file_service().delete_all_conversation_files(
+        session, conversation.id, conversation.file_ids, user_id
+    )
     conversation_crud.delete_conversation(session, conversation_id, user_id)
 
     return DeleteConversationResponse()
@@ -351,7 +352,9 @@ async def upload_file(
     """
 
     user_id = ctx.get_user_id()
-    validate_file_size(session, user_id, file)
+    # Currently do not limit file size for Compass
+    if Settings().feature_flags.use_compass_file_storage is False:
+        validate_file_size(session, user_id, file)
 
     # Create new conversation
     if not conversation_id:
@@ -453,9 +456,14 @@ async def batch_upload_file(
             )
 
     # TODO: check if file already exists in DB once we have files per agents
+
     try:
         uploaded_files = await get_file_service().create_conversation_files(
-            session, files, user_id, conversation.id
+            session,
+            files,
+            user_id,
+            conversation.id,
+            ctx,
         )
     except Exception as e:
         raise HTTPException(
@@ -496,42 +504,6 @@ async def list_files(
     return files_with_conversation_id
 
 
-@router.put("/{conversation_id}/files/{file_id}", response_model=FilePublic)
-async def update_file(
-    conversation_id: str,
-    file_id: str,
-    new_file: UpdateFileRequest,
-    session: DBSessionDep,
-    ctx: Context = Depends(get_context),
-) -> FilePublic:
-    """
-    Update a file by ID.
-
-    Args:
-        conversation_id (str): Conversation ID.
-        file_id (str): File ID.
-        new_file (UpdateFileRequest): New file data.
-        session (DBSessionDep): Database session.
-        ctx (Context): Context object.
-
-    Returns:
-        FilePublic: Updated file.
-
-    Raises:
-        HTTPException: If the conversation with the given ID is not found.
-    """
-    user_id = ctx.get_user_id()
-    _ = validate_conversation(session, conversation_id, user_id)
-    _ = validate_file(session, file_id, user_id)
-
-    file = get_file_service().get_file_by_id(session, file_id, user_id)
-    file = get_file_service().update_file(session, file, new_file)
-    files_with_conversation_id = attach_conversation_id_to_files(
-        conversation_id, [file]
-    )
-    return files_with_conversation_id[0]
-
-
 @router.delete("/{conversation_id}/files/{file_id}")
 async def delete_file(
     conversation_id: str,
@@ -555,12 +527,10 @@ async def delete_file(
     """
     user_id = ctx.get_user_id()
     _ = validate_conversation(session, conversation_id, user_id)
-    _ = validate_file(session, file_id, user_id)
-
-    file = get_file_service().get_file_by_id(session, file_id, user_id)
+    validate_file(session, file_id, user_id, conversation_id)
 
     # Delete the File DB object
-    get_file_service().delete_file_from_conversation(
+    get_file_service().delete_conversation_file_by_id(
         session, conversation_id, file_id, user_id
     )
 

src/backend/schemas/file.py

@@ -10,7 +10,8 @@ class File(BaseModel):
     updated_at: datetime.datetime
 
     user_id: str
-    conversation_id: str
+    conversation_id: Optional[str] = None
+    file_content: Optional[str] = None
     file_name: str
     file_path: str
     file_size: int = Field(default=0, ge=0)