User Management

py/core/__init__.py

@@ -147,6 +147,9 @@
     # LLM provider
     "CompletionConfig",
     "CompletionProvider",
+    # User management provider
+    "UserManagementConfig",
+    "UserManagementProvider",
     ## UTILS
     "RecursiveCharacterTextSplitter",
     "TextSplitter",

py/core/base/__init__.py

@@ -124,6 +124,9 @@
     # LLM provider
     "CompletionConfig",
     "CompletionProvider",
+    # User management provider
+    "UserManagementConfig",
+    "UserManagementProvider",
     ## UTILS
     "RecursiveCharacterTextSplitter",
     "TextSplitter",

py/core/base/api/models/__init__.py

@@ -1,10 +1,8 @@
 from shared.api.models.auth.responses import (
     GenericMessageResponse,
     TokenResponse,
-    UserResponse,
     WrappedGenericMessageResponse,
     WrappedTokenResponse,
-    WrappedUserResponse,
 )
 from shared.api.models.ingestion.responses import (
     CreateVectorIndexResponse,
@@ -46,6 +44,7 @@
     ScoreCompletionResponse,
     ServerStats,
     UserOverviewResponse,
+    UserResponse,
     WrappedAddUserResponse,
     WrappedAnalyticsResponse,
     WrappedAppSettingsResponse,
@@ -63,6 +62,7 @@
     WrappedServerStatsResponse,
     WrappedUserCollectionResponse,
     WrappedUserOverviewResponse,
+    WrappedUserResponse,
     WrappedUsersInCollectionResponse,
     WrappedVerificationResult,
 )

py/core/base/providers/__init__.py

@@ -21,6 +21,7 @@
 from .ingestion import ChunkingStrategy, IngestionConfig, IngestionProvider
 from .llm import CompletionConfig, CompletionProvider
 from .orchestration import OrchestrationConfig, OrchestrationProvider, Workflow
+from .user_management import UserManagementConfig, UserManagementProvider
 
 __all__ = [
     # Auth provider
@@ -64,4 +65,7 @@
     "OrchestrationConfig",
     "OrchestrationProvider",
     "Workflow",
+    # User management provider
+    "UserManagementConfig",
+    "UserManagementProvider",
 ]

py/core/base/providers/database.py

@@ -6,8 +6,6 @@
     Any,
     AsyncGenerator,
     BinaryIO,
-    Dict,
-    List,
     Optional,
     Sequence,
     Tuple,
@@ -50,7 +48,6 @@
     KGEnrichmentEstimationResponse,
     UserResponse,
 )
-from core.base.utils import _decorate_vector_type
 
 from ..logger import RunInfoLog
 from ..logger.base import RunType
@@ -79,15 +76,6 @@
 logger = logging.getLogger()
 
 
-def escape_braces(s: str) -> str:
-    """
-    Escape braces in a string.
-    This is a placeholder function - implement the actual logic as needed.
-    """
-    # Implement your escape_braces logic here
-    return s.replace("{", "{{").replace("}", "}}")
-
-
 logger = logging.getLogger()
 
 
@@ -658,14 +646,14 @@ async def add_triples(
     @abstractmethod
     async def get_entity_map(
         self, offset: int, limit: int, document_id: UUID
-    ) -> Dict[str, Dict[str, List[Dict[str, Any]]]]:
+    ) -> dict[str, dict[str, list[dict[str, Any]]]]:
         """Get entity map for a document."""
         pass
 
     @abstractmethod
     async def upsert_embeddings(
         self,
-        data: List[Tuple[Any]],
+        data: list[Tuple[Any]],
         table_name: str,
     ) -> None:
         """Upsert embeddings into storage."""
@@ -680,7 +668,7 @@ async def vector_query(
 
     # Community management
     @abstractmethod
-    async def add_community_info(self, communities: List[Any]) -> None:
+    async def add_community_info(self, communities: list[Any]) -> None:
         """Add communities to storage."""
         pass
 
@@ -713,22 +701,22 @@ async def get_community_details(
     @abstractmethod
     async def get_community_reports(
         self, collection_id: UUID
-    ) -> List[CommunityReport]:
+    ) -> list[CommunityReport]:
         """Get community reports for a collection."""
         pass
 
     @abstractmethod
     async def check_community_reports_exist(
         self, collection_id: UUID, offset: int, limit: int
-    ) -> List[int]:
+    ) -> list[int]:
         """Check which community reports exist."""
         pass
 
     @abstractmethod
     async def perform_graph_clustering(
         self,
         collection_id: UUID,
-        leiden_params: Dict[str, Any],
+        leiden_params: dict[str, Any],
     ) -> int:
         """Perform graph clustering."""
         pass
@@ -753,10 +741,10 @@ async def delete_node_via_document_id(
     async def get_entities(
         self,
         collection_id: Optional[UUID] = None,
-        entity_ids: Optional[List[str]] = None,
-        entity_names: Optional[List[str]] = None,
+        entity_ids: Optional[list[str]] = None,
+        entity_names: Optional[list[str]] = None,
         entity_table_name: str = "document_entity",
-        extra_columns: Optional[List[str]] = None,
+        extra_columns: Optional[list[str]] = None,
         offset: int = 0,
         limit: int = -1,
     ) -> dict:
@@ -767,8 +755,8 @@ async def get_entities(
     async def get_triples(
         self,
         collection_id: Optional[UUID] = None,
-        entity_names: Optional[List[str]] = None,
-        triple_ids: Optional[List[str]] = None,
+        entity_names: Optional[list[str]] = None,
+        triple_ids: Optional[list[str]] = None,
         offset: int = 0,
         limit: int = -1,
     ) -> dict:
@@ -865,7 +853,7 @@ async def get_existing_entity_extraction_ids(
     @abstractmethod
     async def get_all_triples(
         self, collection_id: UUID, document_ids: Optional[list[UUID]] = None
-    ) -> List[Triple]:
+    ) -> list[Triple]:
         raise NotImplementedError
 
     @abstractmethod
@@ -1002,8 +990,8 @@ async def info_log(
 
     @abstractmethod
     async def get_logs(
-        self, run_ids: List[UUID], limit_per_run: int = 10
-    ) -> List[Dict]:
+        self, run_ids: list[UUID], limit_per_run: int = 10
+    ) -> list[dict]:
         """Retrieve logs for specified run IDs."""
         pass
 
@@ -1013,8 +1001,8 @@ async def get_info_logs(
         offset: int = 0,
         limit: int = 100,
         run_type_filter: Optional[RunType] = None,
-        user_ids: Optional[List[UUID]] = None,
-    ) -> List[RunInfoLog]:
+        user_ids: Optional[list[UUID]] = None,
+    ) -> list[RunInfoLog]:
         """Retrieve run information logs with filtering options."""
         pass
 
@@ -1032,10 +1020,10 @@ async def delete_conversation(self, conversation_id: str) -> None:
     @abstractmethod
     async def get_conversations_overview(
         self,
-        conversation_ids: Optional[List[UUID]] = None,
+        conversation_ids: Optional[list[UUID]] = None,
         offset: int = 0,
         limit: int = -1,
-    ) -> Dict[str, Union[List[Dict], int]]:
+    ) -> dict[str, Union[list[dict], int]]:
         """Get an overview of conversations with pagination."""
         pass
 
@@ -1046,7 +1034,7 @@ async def add_message(
         conversation_id: str,
         content: Message,
         parent_id: Optional[str] = None,
-        metadata: Optional[Dict] = None,
+        metadata: Optional[dict] = None,
     ) -> str:
         """Add a message to a conversation."""
         pass
@@ -1061,13 +1049,13 @@ async def edit_message(
     @abstractmethod
     async def get_conversation(
         self, conversation_id: str, branch_id: Optional[str] = None
-    ) -> List[Tuple[str, Message]]:
+    ) -> list[Tuple[str, Message]]:
         """Retrieve all messages in a conversation branch."""
         pass
 
     # Branch management methods
     @abstractmethod
-    async def get_branches_overview(self, conversation_id: str) -> List[Dict]:
+    async def get_branches_overview(self, conversation_id: str) -> list[dict]:
         """Get an overview of all branches in a conversation."""
         pass
 
@@ -1556,20 +1544,20 @@ async def add_triples(
 
     async def get_entity_map(
         self, offset: int, limit: int, document_id: UUID
-    ) -> Dict[str, Dict[str, List[Dict[str, Any]]]]:
+    ) -> dict[str, dict[str, list[dict[str, Any]]]]:
         """Forward to KG handler get_entity_map method."""
         return await self.kg_handler.get_entity_map(offset, limit, document_id)
 
     async def upsert_embeddings(
         self,
-        data: List[Tuple[Any]],
+        data: list[Tuple[Any]],
         table_name: str,
     ) -> None:
         """Forward to KG handler upsert_embeddings method."""
         return await self.kg_handler.upsert_embeddings(data, table_name)
 
     # Community methods
-    async def add_community_info(self, communities: List[Any]) -> None:
+    async def add_community_info(self, communities: list[Any]) -> None:
         """Forward to KG handler add_communities method."""
         return await self.kg_handler.add_community_info(communities)
 
@@ -1606,13 +1594,13 @@ async def get_community_details(
 
     async def get_community_reports(
         self, collection_id: UUID
-    ) -> List[CommunityReport]:
+    ) -> list[CommunityReport]:
         """Forward to KG handler get_community_reports method."""
         return await self.kg_handler.get_community_reports(collection_id)
 
     async def check_community_reports_exist(
         self, collection_id: UUID, offset: int, limit: int
-    ) -> List[int]:
+    ) -> list[int]:
         """Forward to KG handler check_community_reports_exist method."""
         return await self.kg_handler.check_community_reports_exist(
             collection_id, offset, limit
@@ -1621,7 +1609,7 @@ async def check_community_reports_exist(
     async def perform_graph_clustering(
         self,
         collection_id: UUID,
-        leiden_params: Dict[str, Any],
+        leiden_params: dict[str, Any],
     ) -> int:
         """Forward to KG handler perform_graph_clustering method."""
         return await self.kg_handler.perform_graph_clustering(
@@ -1649,10 +1637,10 @@ async def delete_node_via_document_id(
     async def get_entities(
         self,
         collection_id: Optional[UUID],
-        entity_ids: Optional[List[str]] = None,
-        entity_names: Optional[List[str]] = None,
+        entity_ids: Optional[list[str]] = None,
+        entity_names: Optional[list[str]] = None,
         entity_table_name: str = "document_entity",
-        extra_columns: Optional[List[str]] = None,
+        extra_columns: Optional[list[str]] = None,
         offset: int = 0,
         limit: int = -1,
     ) -> dict:
@@ -1670,8 +1658,8 @@ async def get_entities(
     async def get_triples(
         self,
         collection_id: Optional[UUID] = None,
-        entity_names: Optional[List[str]] = None,
-        triple_ids: Optional[List[str]] = None,
+        entity_names: Optional[list[str]] = None,
+        triple_ids: Optional[list[str]] = None,
         offset: int = 0,
         limit: int = -1,
     ) -> dict:
@@ -1735,7 +1723,7 @@ async def get_deduplication_estimate(
 
     async def get_all_triples(
         self, collection_id: UUID, document_ids: Optional[list[UUID]] = None
-    ) -> List[Triple]:
+    ) -> list[Triple]:
         return await self.kg_handler.get_all_triples(
             collection_id, document_ids
         )
@@ -1874,18 +1862,18 @@ async def get_info_logs(
         offset: int = 0,
         limit: int = 100,
         run_type_filter: Optional[RunType] = None,
-        user_ids: Optional[List[UUID]] = None,
-    ) -> List[RunInfoLog]:
+        user_ids: Optional[list[UUID]] = None,
+    ) -> list[RunInfoLog]:
         """Retrieve log info entries with filtering and pagination."""
         return await self.logging_handler.get_info_logs(
             offset, limit, run_type_filter, user_ids
         )
 
     async def get_logs(
         self,
-        run_ids: List[UUID],
+        run_ids: list[UUID],
         limit_per_run: int = 10,
-    ) -> List[Dict[str, Any]]:
+    ) -> list[dict[str, Any]]:
         """Retrieve logs for specified run IDs with a per-run limit."""
         return await self.logging_handler.get_logs(run_ids, limit_per_run)
 
@@ -1899,10 +1887,10 @@ async def delete_conversation(self, conversation_id: str) -> None:
 
     async def get_conversations_overview(
         self,
-        conversation_ids: Optional[List[UUID]] = None,
+        conversation_ids: Optional[list[UUID]] = None,
         offset: int = 0,
         limit: int = -1,
-    ) -> Dict[str, Union[List[Dict], int]]:
+    ) -> dict[str, Union[list[dict], int]]:
         """Get an overview of conversations with pagination."""
         return await self.logging_handler.get_conversations_overview(
             conversation_ids, offset, limit
@@ -1913,7 +1901,7 @@ async def add_message(
         conversation_id: str,
         content: Message,
         parent_id: Optional[str] = None,
-        metadata: Optional[Dict] = None,
+        metadata: Optional[dict] = None,
     ) -> str:
         """Add a message to a conversation."""
         return await self.logging_handler.add_message(
@@ -1928,13 +1916,13 @@ async def edit_message(
 
     async def get_conversation(
         self, conversation_id: str, branch_id: Optional[str] = None
-    ) -> List[Tuple[str, Message]]:
+    ) -> list[Tuple[str, Message]]:
         """Retrieve all messages in a conversation branch."""
         return await self.logging_handler.get_conversation(
             conversation_id, branch_id
         )
 
-    async def get_branches_overview(self, conversation_id: str) -> List[Dict]:
+    async def get_branches_overview(self, conversation_id: str) -> list[dict]:
         """Get an overview of all branches in a conversation."""
         return await self.logging_handler.get_branches_overview(
             conversation_id