feat: kms file size limit

backend/api/quivr_api/models/settings.py

@@ -103,6 +103,8 @@ def set_once_user_properties(self, user_id: UUID, event_name, properties: dict):
 
 class BrainSettings(BaseSettings):
     model_config = SettingsConfigDict(validate_default=False)
+    pg_database_url: str
+    pg_database_async_url: str
     openai_api_key: str = ""
     azure_openai_embeddings_url: str = ""
     supabase_url: str = ""
@@ -112,11 +114,10 @@ class BrainSettings(BaseSettings):
     ollama_api_base_url: str | None = None
     langfuse_public_key: str | None = None
     langfuse_secret_key: str | None = None
-    pg_database_url: str
-    pg_database_async_url: str
     sqlalchemy_pool_size: int = 5
     sqlalchemy_max_pool_overflow: int = 5
     embedding_dim: int = 1536
+    max_file_size: int = int(5e7)
 
 
 class ResendSettings(BaseSettings):

backend/api/quivr_api/modules/knowledge/controller/knowledge_routes.py

@@ -18,6 +18,7 @@
 from quivr_api.celery_config import celery
 from quivr_api.logger import get_logger
 from quivr_api.middlewares.auth import AuthBearer, get_current_user
+from quivr_api.models.settings import settings
 from quivr_api.modules.brain.service.brain_authorization_service import (
     validate_brain_authorization,
 )
@@ -123,6 +124,11 @@ async def create_knowledge(
     current_user: UserIdentity = Depends(get_current_user),
 ):
     knowledge = AddKnowledge.model_validate_json(knowledge_data)
+    if file and file.size and file.size > settings.max_file_size:
+        raise HTTPException(
+            status_code=status.HTTP_413_REQUEST_ENTITY_TOO_LARGE,
+            detail="Uploaded file is too large",
+        )
     if not knowledge.file_name and not knowledge.url:
         raise HTTPException(
             status_code=status.HTTP_400_BAD_REQUEST,

backend/api/quivr_api/modules/knowledge/tests/test_knowledge_controller.py

@@ -5,6 +5,7 @@
 import pytest
 import pytest_asyncio
 from httpx import ASGITransport, AsyncClient
+from quivr_api.models.settings import BrainSettings
 from quivr_core.models import KnowledgeStatus
 from sqlmodel import select
 from sqlmodel.ext.asyncio.session import AsyncSession
@@ -144,6 +145,33 @@ async def test_post_knowledge_folder(test_client: AsyncClient):
     assert km.children == []
 
 
+@pytest.mark.asyncio(loop_scope="session")
+async def test_add_knowledge_large_file(monkeypatch, test_client):
+    _settings = BrainSettings()
+    _settings.max_file_size = 2
+    monkeypatch.setattr(
+        "quivr_api.modules.knowledge.controller.knowledge_routes.settings", _settings
+    )
+    km_data = {
+        "file_name": "test_file.txt",
+        "source": "local",
+        "is_folder": False,
+        "parent_id": None,
+    }
+
+    multipart_data = {
+        "knowledge_data": (None, json.dumps(km_data), "application/json"),
+        "file": ("test_file.txt", b"Test file content", "application/octet-stream"),
+    }
+
+    response = await test_client.post(
+        "/knowledge/",
+        files=multipart_data,
+    )
+
+    assert response.status_code == 413
+
+
 @pytest.mark.asyncio(loop_scope="session")
 async def test_add_knowledge_invalid_input(test_client):
     response = await test_client.post("/knowledge/", files={})

backend/requirements-dev.lock

@@ -553,6 +553,7 @@ llama-parse==0.5.6
     # via quivr-api
 llvmlite==0.43.0
     # via numba
+locust==2.31.8
 lxml==5.3.0
     # via pikepdf
     # via python-docx
@@ -1130,6 +1131,8 @@ sentry-sdk==2.13.0
     # via quivr-api
 setuptools==70.0.0
     # via opentelemetry-instrumentation
+    # via zope-event
+    # via zope-interface
 shapely==2.0.6
     # via python-doctr
 simple-websocket==1.0.0