Merge pull request #3272 from QuivrHQ/release/load-test-kms

feat: load test KMS

.env.example

@@ -54,6 +54,8 @@ EXTERNAL_SUPABASE_URL=http://localhost:54321
 SUPABASE_SERVICE_KEY=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJzdXBhYmFzZS1kZW1vIiwicm9sZSI6InNlcnZpY2Vfcm9sZSIsImV4cCI6MTk4MzgxMjk5Nn0.EGIM96RAZx35lJzdJsyH-qQwv8Hdp7fsn3W0YpN81IU
 PG_DATABASE_URL=postgresql://postgres:[email protected]:54322/postgres
 PG_DATABASE_ASYNC_URL=postgresql+asyncpg://postgres:[email protected]:54322/postgres
+SQLALCHEMY_POOL_SIZE=10
+SQLALCHEMY_MAX_POOL_OVERFLOW=0
 JWT_SECRET_KEY=super-secret-jwt-token-with-at-least-32-characters-long
 AUTHENTICATE=true
 TELEMETRY_ENABLED=true

.gitignore

@@ -103,3 +103,4 @@ backend/core/examples/chatbot/.chainlit/translations/en-US.json
 .tox
 Pipfile
 *.pkl
+backend/benchmarks/data.json

backend/api/quivr_api/models/settings.py

@@ -123,6 +123,8 @@ class BrainSettings(BaseSettings):
     pg_database_url: str
     pg_database_async_url: str
     embedding_dim: int
+    sqlalchemy_pool_size: int
+    sqlalchemy_max_pool_overflow: int
 
 
 class ResendSettings(BaseSettings):

backend/api/quivr_api/modules/dependencies.py

@@ -61,16 +61,20 @@ def get_repository_cls(cls) -> Type[R]:
     future=True,
     # NOTE: pessimistic bound on
     pool_pre_ping=True,
-    pool_size=10,  # NOTE: no bouncer for now, if 6 process workers => 6
+    pool_size=1,  # NOTE: no bouncer for now, if 6 process workers => 6
+    max_overflow=0,
     pool_recycle=1800,
 )
 async_engine = create_async_engine(
     settings.pg_database_async_url,
-    connect_args={"server_settings": {"application_name": "quivr-api-async"}},
+    connect_args={
+        "server_settings": {"application_name": "quivr-api-async"},
+    },
     echo=True if os.getenv("ORM_DEBUG") else False,
     future=True,
     pool_pre_ping=True,
-    pool_size=5,  # NOTE: no bouncer for now, if 6 process workers => 6
+    pool_size=settings.sqlalchemy_pool_size,  # NOTE: no bouncer for now, if 6 process workers => 6
+    max_overflow=settings.sqlalchemy_max_pool_overflow,
     pool_recycle=1800,
     isolation_level="AUTOCOMMIT",
 )

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

@@ -3,15 +3,22 @@
 from typing import List, Optional
 from uuid import UUID
 
-from fastapi import APIRouter, Depends, File, HTTPException, Query, UploadFile, status
+from fastapi import (
+    APIRouter,
+    Depends,
+    File,
+    HTTPException,
+    Query,
+    Response,
+    UploadFile,
+    status,
+)
 from quivr_core.models import KnowledgeStatus
 
 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.modules.brain.entity.brain_entity import RoleEnum
 from quivr_api.modules.brain.service.brain_authorization_service import (
-    has_brain_authorization,
     validate_brain_authorization,
 )
 from quivr_api.modules.dependencies import get_service
@@ -226,33 +233,6 @@ async def update_knowledge(
         raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR)
 
 
-@knowledge_router.delete(
-    "/knowledge/{knowledge_id}",
-    dependencies=[
-        Depends(AuthBearer()),
-        Depends(has_brain_authorization(RoleEnum.Owner)),
-    ],
-    tags=["Knowledge"],
-)
-async def delete_knowledge_brain(
-    knowledge_id: UUID,
-    knowledge_service: KnowledgeService = Depends(get_knowledge_service),
-    current_user: UserIdentity = Depends(get_current_user),
-    brain_id: UUID = Query(..., description="The ID of the brain"),
-):
-    """
-    Delete a specific knowledge from a brain.
-    """
-
-    knowledge = await knowledge_service.get_knowledge(knowledge_id)
-    file_name = knowledge.file_name if knowledge.file_name else knowledge.url
-    await knowledge_service.remove_knowledge_brain(brain_id, knowledge_id)
-
-    return {
-        "message": f"{file_name} of brain {brain_id} has been deleted by user {current_user.email}."
-    }
-
-
 @knowledge_router.delete(
     "/knowledge/{knowledge_id}",
     status_code=status.HTTP_202_ACCEPTED,
@@ -297,7 +277,7 @@ async def link_knowledge_to_brain(
         link_request.bulk_id,
     )
     if len(brains_ids) == 0:
-        return "empty brain list"
+        return Response(status_code=status.HTTP_204_NO_CONTENT)
 
     if knowledge_dto.id is None:
         if knowledge_dto.sync_file_id is None:
@@ -319,25 +299,25 @@ async def link_knowledge_to_brain(
             knowledge_dto.id, brains_ids=brains_ids, user_id=current_user.id
         )
 
-    for knowledge in filter(
-        lambda k: k.status
-        not in [KnowledgeStatus.PROCESSED, KnowledgeStatus.PROCESSING],
-        linked_kms,
-    ):
-        assert knowledge.id
+    for knowledge in [
+        k
+        for k in linked_kms
+        if await k.awaitable_attrs.status
+        not in [KnowledgeStatus.PROCESSED, KnowledgeStatus.PROCESSING]
+    ]:
         upload_notification = notification_service.add_notification(
             CreateNotification(
                 user_id=current_user.id,
                 bulk_id=bulk_id,
                 status=NotificationsStatusEnum.INFO,
-                title=f"{knowledge.file_name}",
+                title=f"{await knowledge.awaitable_attrs.file_name}",
                 category="process",
             )
         )
         celery.send_task(
             "process_file_task",
             kwargs={
-                "knowledge_id": knowledge.id,
+                "knowledge_id": await knowledge.awaitable_attrs.id,
                 "notification_id": upload_notification.id,
             },
         )

backend/api/quivr_api/modules/knowledge/entity/knowledge.py

@@ -99,10 +99,10 @@ async def to_dto(
         self, get_children: bool = True, get_parent: bool = True
     ) -> KnowledgeDTO:
         assert (
-            self.updated_at
+            await self.awaitable_attrs.updated_at
         ), "knowledge should be inserted before transforming to dto"
         assert (
-            self.created_at
+            await self.awaitable_attrs.created_at
         ), "knowledge should be inserted before transforming to dto"
         brains = await self.awaitable_attrs.brains
         children: list[KnowledgeDB] = (
@@ -125,8 +125,8 @@ async def to_dto(
             is_folder=self.is_folder,
             file_size=self.file_size or 0,
             file_sha1=self.file_sha1,
-            updated_at=self.updated_at,
-            created_at=self.created_at,
+            updated_at=await self.awaitable_attrs.updated_at,
+            created_at=await self.awaitable_attrs.created_at,
             metadata=self.metadata_,  # type: ignore
             brains=[b.model_dump() for b in brains],
             parent=parent_dto,

backend/api/quivr_api/modules/knowledge/service/knowledge_service.py

@@ -142,9 +142,11 @@ async def get_knowledge(
 
     async def update_knowledge(
         self,
-        knowledge: KnowledgeDB,
+        knowledge: KnowledgeDB | UUID,
         payload: KnowledgeDTO | KnowledgeUpdate | dict[str, Any],
     ):
+        if isinstance(knowledge, UUID):
+            knowledge = await self.repository.get_knowledge_by_id(knowledge)
         return await self.repository.update_knowledge(knowledge, payload)
 
     async def create_knowledge(

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

@@ -96,16 +96,17 @@ async def test_service():
 
 
 @pytest.mark.asyncio(loop_scope="session")
-async def test_post_knowledge_folder(test_client: AsyncClient):
+async def test_post_knowledge(test_client: AsyncClient):
     km_data = {
         "file_name": "test_file.txt",
         "source": "local",
-        "is_folder": True,
+        "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(
@@ -114,26 +115,19 @@ async def test_post_knowledge_folder(test_client: AsyncClient):
     )
 
     assert response.status_code == 200
-    km = KnowledgeDTO.model_validate(response.json())
-
-    assert km.id
-    assert km.is_folder
-    assert km.parent is None
-    assert km.children == []
 
 
 @pytest.mark.asyncio(loop_scope="session")
-async def test_post_knowledge(test_client: AsyncClient):
+async def test_post_knowledge_folder(test_client: AsyncClient):
     km_data = {
         "file_name": "test_file.txt",
         "source": "local",
-        "is_folder": False,
+        "is_folder": True,
         "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(
@@ -142,6 +136,12 @@ async def test_post_knowledge(test_client: AsyncClient):
     )
 
     assert response.status_code == 200
+    km = KnowledgeDTO.model_validate(response.json())
+
+    assert km.id
+    assert km.is_folder
+    assert km.parent is None
+    assert km.children == []
 
 
 @pytest.mark.asyncio(loop_scope="session")

backend/benchmarks/benchmark_kms.sh

@@ -0,0 +1,37 @@
+#!/bin/bash
+
+# Function to handle cleanup on exit
+cleanup() {
+    echo "Cleaning up..."
+    # Stop Uvicorn server if running
+    if [[ ! -z "$UVICORN_PID" ]]; then
+        kill "$UVICORN_PID"
+        wait "$UVICORN_PID"
+    fi
+    exit 0
+}
+
+# Trap signals (like Ctrl+C, SIGTERM) to run the cleanup function
+#trap cleanup SIGINT SIGTERM
+
+# Reset and start Supabase
+supabase db reset && supabase stop && supabase start
+
+# Remove old benchmark data
+rm -f benchmarks/data.json
+
+# Load new data
+rye run python benchmarks/load_data.py
+
+# Start Uvicorn server in the background
+LOG_LEVEL=info rye run uvicorn quivr_api.main:app --log-level info --host 0.0.0.0 --port 5050 --workers 5 --loop uvloop &
+UVICORN_PID=$!
+
+# Wait a bit to ensure the server is running
+sleep 1
+
+# Run Locust for benchmarking
+rye run locust -f benchmarks/locustfile_kms.py -H http://localhost:5050
+
+# Wait for all background processes (including Uvicorn) to finish
+wait "$UVICORN_PID"