release: kms 0.1

.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

.vscode/settings.json

@@ -1,8 +1,9 @@
 {
   "editor.codeActionsOnSave": {
     "source.organizeImports": "explicit",
+    "source.organizeImports.ruff": "explicit",
     "source.fixAll": "explicit",
-    "source.unusedImports": "explicit",
+    "source.unusedImports": "explicit"
   },
   "editor.formatOnSave": true,
   "editor.formatOnSaveMode": "file",
@@ -24,9 +25,7 @@
       "source.fixAll": "explicit"
     }
   },
-  "python.analysis.extraPaths": [
-    "./backend"
-  ],
+  "python.analysis.extraPaths": ["./backend"],
   "python.defaultInterpreterPath": "python3",
   "python.testing.pytestArgs": [
     "-v",
@@ -43,6 +42,5 @@
     "reportMissingImports": "error",
     "reportUnusedImport": "warning",
     "reportGeneralTypeIssues": "warning"
-  },
-  "makefile.configureOnOpen": false
+  }
 }

backend/api/quivr_api/__init__.py

@@ -1,15 +1,18 @@
 from quivr_api.modules.brain.entity.brain_entity import Brain
+from quivr_api.modules.brain.entity.brain_user import BrainUserDB
 from quivr_api.modules.knowledge.entity.knowledge import KnowledgeDB
 
 from .modules.chat.entity.chat import Chat, ChatHistory
-from .modules.sync.entity.sync_models import NotionSyncFile
+from .modules.sync.entity.sync_models import NotionSyncFile, Sync
 from .modules.user.entity.user_identity import User
 
 __all__ = [
     "Chat",
     "ChatHistory",
+    "BrainUserDB",
     "User",
     "NotionSyncFile",
     "KnowledgeDB",
     "Brain",
+    "Sync",
 ]

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,9 +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/brain/entity/brain_entity.py

@@ -10,6 +10,7 @@
 from sqlmodel import TIMESTAMP, Column, Field, Relationship, SQLModel, text
 from sqlmodel import UUID as PGUUID
 
+from quivr_api.modules.brain.entity.brain_user import BrainUserDB
 from quivr_api.modules.brain.entity.integration_brain import (
     IntegrationDescriptionEntity,
     IntegrationEntity,
@@ -69,7 +70,12 @@ class Brain(AsyncAttrs, SQLModel, table=True):
     knowledges: List[KnowledgeDB] = Relationship(
         back_populates="brains", link_model=KnowledgeBrain
     )
-
+    users: List["User"] = Relationship(  # type: ignore # noqa: F821
+        back_populates="brains",
+        link_model=BrainUserDB,
+    )
+    snippet_color: str | None = Field(default="#d0c6f2")
+    snippet_emoji: str | None = Field(default="🧠")
     # TODO : add
     # "meaning" "public"."vector",
     # "tags" "public"."tags"[]

backend/api/quivr_api/modules/brain/entity/brain_user.py

@@ -0,0 +1,14 @@
+from uuid import UUID
+
+from sqlmodel import Field, SQLModel
+
+
+class BrainUserDB(SQLModel, table=True):
+    __tablename__ = "brains_users"  # type: ignore
+
+    brain_id: UUID = Field(
+        nullable=False, foreign_key="brains.brain_id", primary_key=True
+    )
+    user_id: UUID = Field(nullable=False, foreign_key="users.id", primary_key=True)
+    default_brain: bool
+    rights: str

backend/api/quivr_api/modules/brain/repository/brains_users.py

@@ -2,7 +2,7 @@
 
 from quivr_api.logger import get_logger
 from quivr_api.modules.brain.entity.brain_entity import (
-    BrainUser,
+    BrainUserDB,
     MinimalUserBrainEntity,
 )
 from quivr_api.modules.brain.repository.interfaces.brains_users_interface import (
@@ -161,15 +161,15 @@ def get_user_default_brain_id(self, user_id: UUID) -> UUID | None:
             return None
         return UUID(response[0].get("brain_id"))
 
-    def get_brain_users(self, brain_id: UUID) -> list[BrainUser]:
+    def get_brain_users(self, brain_id: UUID) -> list[BrainUserDB]:
         response = (
             self.db.table("brains_users")
             .select("id:brain_id, *")
             .filter("brain_id", "eq", str(brain_id))
             .execute()
         )
 
-        return [BrainUser(**item) for item in response.data]
+        return [BrainUserDB(**item) for item in response.data]
 
     def delete_brain_subscribers(self, brain_id: UUID):
         results = (

backend/api/quivr_api/modules/brain/repository/interfaces/brains_users_interface.py

@@ -3,7 +3,7 @@
 from uuid import UUID
 
 from quivr_api.modules.brain.entity.brain_entity import (
-    BrainUser,
+    BrainUserDB,
     MinimalUserBrainEntity,
 )
 
@@ -56,7 +56,7 @@ def get_user_default_brain_id(self, user_id: UUID) -> UUID | None:
         pass
 
     @abstractmethod
-    def get_brain_users(self, brain_id: UUID) -> List[BrainUser]:
+    def get_brain_users(self, brain_id: UUID) -> List[BrainUserDB]:
         """
         Get all users for a brain
         """
@@ -88,7 +88,7 @@ def update_brain_user_default_status(
     @abstractmethod
     def update_brain_user_rights(
         self, brain_id: UUID, user_id: UUID, rights: str
-    ) -> BrainUser:
+    ) -> BrainUserDB:
         """
         Update the rights for a user in a brain
         """

backend/api/quivr_api/modules/brain/service/brain_service.py

@@ -1,4 +1,4 @@
-from typing import Dict, Optional, Tuple
+from typing import Optional, Tuple
 from uuid import UUID
 
 from fastapi import HTTPException
@@ -54,7 +54,7 @@ def find_brain_from_question(
         chat_id: UUID,
         history,
         vector_store: CustomSupabaseVectorStore,
-    ) -> Tuple[Optional[BrainEntity], Dict[str, str]]:
+    ) -> Tuple[Optional[BrainEntity], dict[str, str]]:
         """Find the brain to use for a question.
 
         Args:

backend/api/quivr_api/modules/brain/service/brain_user_service.py

@@ -6,7 +6,7 @@
 from quivr_api.logger import get_logger
 from quivr_api.modules.brain.entity.brain_entity import (
     BrainEntity,
-    BrainUser,
+    BrainUserDB,
     MinimalUserBrainEntity,
     RoleEnum,
 )
@@ -74,7 +74,7 @@ def get_user_brains(self, user_id: UUID) -> list[MinimalUserBrainEntity]:
 
         return results  # type: ignore
 
-    def get_brain_users(self, brain_id: UUID) -> List[BrainUser]:
+    def get_brain_users(self, brain_id: UUID) -> List[BrainUserDB]:
         return self.brain_user_repository.get_brain_users(brain_id)
 
     def update_brain_user_rights(

backend/api/quivr_api/modules/dependencies.py

@@ -5,12 +5,7 @@
 from fastapi import Depends
 from langchain.embeddings.base import Embeddings
 from langchain_community.embeddings.ollama import OllamaEmbeddings
-
-# from langchain_community.vectorstores.supabase import SupabaseVectorStore
 from langchain_openai import AzureOpenAIEmbeddings, OpenAIEmbeddings
-
-# from quivr_api.modules.vector.service.vector_service import VectorService
-# from quivr_api.modules.vectorstore.supabase import CustomSupabaseVectorStore
 from sqlalchemy import Engine, create_engine
 from sqlalchemy.ext.asyncio import create_async_engine
 from sqlmodel import Session, text
@@ -61,16 +56,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",
 )