Clean up type errors, cleaner code

js/sdk/src/v3/clients/graphs.ts

@@ -565,9 +565,34 @@ export class GraphsClient {
     );
   }
 
+  /**
+   * Creates communities in the graph by analyzing entity relationships and similarities.
+   *
+   * Communities are created through the following process:
+   * 1. Analyzes entity relationships and metadata to build a similarity graph
+   * 2. Applies advanced community detection algorithms (e.g. Leiden) to identify densely connected groups
+   * 3. Creates hierarchical community structure with multiple granularity levels
+   * 4. Generates natural language summaries and statistical insights for each community
+   *
+   * The resulting communities can be used to:
+   * - Understand high-level graph structure and organization
+   * - Identify key entity groupings and their relationships
+   * - Navigate and explore the graph at different levels of detail
+   * - Generate insights about entity clusters and their characteristics
+   *
+   * The community detection process is configurable through settings like:
+   * - Community detection algorithm parameters
+   * - Summary generation prompt
+   *
+   * @param options
+   * @returns
+   */
   @feature("graphs.buildCommunities")
   async buildCommunities(options: {
     collectionId: string;
+    runType?: string;
+    kgEntichmentSettings?: Record<string, any>;
+    runWithOrchestration?: boolean;
   }): Promise<WrappedBooleanResponse> {
     return this.client.makeRequest(
       "POST",

py/core/main/api/v3/graph_router.py

@@ -304,7 +304,8 @@ async def build_communities(
             run_with_orchestration: Optional[bool] = Body(True),
             auth_user=Depends(self.providers.auth.auth_wrapper),
         ):  # -> WrappedKGEnrichmentResponse:
-            """Creates communities in the graph by analyzing entity relationships and similarities.
+            """
+            Creates communities in the graph by analyzing entity relationships and similarities.
 
             Communities are created through the following process:
             1. Analyzes entity relationships and metadata to build a similarity graph
@@ -323,8 +324,14 @@ async def build_communities(
                 - Summary generation prompt
             """
             print("collection_id = ", collection_id)
-            if not auth_user.is_superuser:
-                logger.warning("Implement permission checks here.")
+            if (
+                not auth_user.is_superuser
+                and collection_id not in auth_user.graph_ids
+            ):
+                raise R2RException(
+                    "The currently authenticated user does not have access to the specified graph.",
+                    403,
+                )
 
             # If no collection ID is provided, use the default user collection
             # id = generate_default_user_collection_id(auth_user.id)
@@ -586,6 +593,14 @@ async def get_entities(
             auth_user=Depends(self.providers.auth.auth_wrapper),
         ) -> WrappedEntitiesResponse:
             """Lists all entities in the graph with pagination support."""
+            if (
+                not auth_user.is_superuser
+                and collection_id not in auth_user.graph_ids
+            ):
+                raise R2RException(
+                    "The currently authenticated user does not have access to the specified graph.",
+                    403,
+                )
             # return await self.services["kg"].get_entities(
             #     id, offset, limit, auth_user
             # )
@@ -626,7 +641,7 @@ async def create_entity(
                 and collection_id not in auth_user.graph_ids
             ):
                 raise R2RException(
-                    "The currently authenticated user does not have access to this graph.",
+                    "The currently authenticated user does not have access to the specified graph.",
                     403,
                 )
 
@@ -680,7 +695,7 @@ async def create_relationship(
                 and collection_id not in auth_user.graph_ids
             ):
                 raise R2RException(
-                    "The currently authenticated user does not have access to this graph.",
+                    "The currently authenticated user does not have access to the specified graph.",
                     403,
                 )
 
@@ -750,7 +765,15 @@ async def get_entity(
             auth_user=Depends(self.providers.auth.auth_wrapper),
         ) -> WrappedEntityResponse:
             """Retrieves a specific entity by its ID."""
-            # Note: The original was missing implementation, so assuming similar pattern to relationships
+            if (
+                not auth_user.is_superuser
+                and collection_id not in auth_user.graph_ids
+            ):
+                raise R2RException(
+                    "The currently authenticated user does not have access to the specified graph.",
+                    403,
+                )
+
             result = await self.providers.database.graph_handler.entities.get(
                 collection_id, "graph", entity_ids=[entity_id]
             )
@@ -856,6 +879,15 @@ async def delete_entity(
             auth_user=Depends(self.providers.auth.auth_wrapper),
         ) -> WrappedBooleanResponse:
             """Removes an entity from the graph."""
+            if (
+                not auth_user.is_superuser
+                and collection_id not in auth_user.graph_ids
+            ):
+                raise R2RException(
+                    "The currently authenticated user does not have access to the specified graph.",
+                    403,
+                )
+
             await self.providers.database.graph_handler.entities.delete(
                 collection_id, [entity_id], "graph"
             )
@@ -922,13 +954,12 @@ async def get_relationships(
             """
             Lists all relationships in the graph with pagination support.
             """
-            # Permission check
             if (
                 not auth_user.is_superuser
                 and collection_id not in auth_user.graph_ids
             ):
                 raise R2RException(
-                    "The currently authenticated user does not have access to this graph.",
+                    "The currently authenticated user does not have access to the specified graph.",
                     403,
                 )
 
@@ -1000,6 +1031,15 @@ async def get_relationship(
             auth_user=Depends(self.providers.auth.auth_wrapper),
         ) -> WrappedRelationshipResponse:
             """Retrieves a specific relationship by its ID."""
+            if (
+                not auth_user.is_superuser
+                and collection_id not in auth_user.graph_ids
+            ):
+                raise R2RException(
+                    "The currently authenticated user does not have access to the specified graph.",
+                    403,
+                )
+
             results = (
                 await self.providers.database.graph_handler.relationships.get(
                     collection_id, "graph", relationship_ids=[relationship_id]
@@ -1126,6 +1166,15 @@ async def delete_relationship(
             auth_user=Depends(self.providers.auth.auth_wrapper),
         ) -> WrappedBooleanResponse:
             """Removes a relationship from the graph."""
+            if (
+                not auth_user.is_superuser
+                and collection_id not in auth_user.graph_ids
+            ):
+                raise R2RException(
+                    "The currently authenticated user does not have access to the specified graph.",
+                    403,
+                )
+
             # return await self.services[
             #     "kg"
             # ].documents.graph_handler.relationships.remove_from_graph(
@@ -1632,6 +1681,7 @@ async def pull(
                     "The currently authenticated user does not have access to the specified graph.",
                     403,
                 )
+
             list_graphs_response = await self.services["kg"].list_graphs(
                 # user_ids=None,
                 graph_ids=[collection_id],
@@ -1766,7 +1816,6 @@ async def remove_document(
 
             The user must have access to both the graph and the document being removed.
             """
-            # Check user permissions for graph
             if (
                 not auth_user.is_superuser
                 and collection_id not in auth_user.graph_ids
@@ -1776,7 +1825,6 @@ async def remove_document(
                     403,
                 )
 
-            # Check user permissions for document
             if (
                 not auth_user.is_superuser
                 and document_id not in auth_user.document_ids

py/core/main/services/kg_service.py

@@ -207,32 +207,18 @@ async def update_entity(
         )
 
     @telemetry_event("delete_entity")
-    async def delete_entity_v3(
+    async def delete_entity(
         self,
         id: UUID,
         entity_id: UUID,
         level: DataLevel,
-        **kwargs,
     ):
         return await self.providers.database.graph_handler.entities.delete(
             id=id,
             entity_id=entity_id,
             level=level,
         )
 
-    @telemetry_event("add_entity_to_graph")
-    async def add_entity_to_graph(
-        self,
-        graph_id: UUID,
-        entity_id: UUID,
-        auth_user: Optional[Any] = None,
-    ):
-        return (
-            await self.providers.database.graph_handler.entities.add_to_graph(
-                graph_id, entity_id, auth_user
-            )
-        )
-
     # TODO: deprecate this
     @telemetry_event("get_entities")
     async def get_entities(
@@ -312,8 +298,8 @@ async def create_relationship(
             )
         )
 
-    @telemetry_event("delete_relationship_v3")
-    async def delete_relationship_v3(
+    @telemetry_event("delete_relationship")
+    async def delete_relationship(
         self,
         id: UUID,
         relationship_id: UUID,
@@ -365,19 +351,19 @@ async def update_relationship(
     @telemetry_event("get_triples")
     async def get_relationships(
         self,
-        collection_id: Optional[UUID] = None,
+        offset: int,
+        limit: int,
+        collection_id: UUID,
         entity_names: Optional[list[str]] = None,
-        relationship_ids: Optional[list[str]] = None,
-        offset: Optional[int] = None,
-        limit: Optional[int] = None,
-        **kwargs,
+        relationship_ids: Optional[list[UUID]] = None,
     ):
-        return await self.providers.database.graph_handler.get_relationships(
-            collection_id=collection_id,
+        return await self.providers.database.graph_handler.relationships.get(
+            parent_id=collection_id,
+            store_type="graph",  # type: ignore
             entity_names=entity_names,
             relationship_ids=relationship_ids,
-            offset=offset or 0,
-            limit=limit or -1,
+            offset=offset,
+            limit=limit,
         )
 
     ################### COMMUNITIES ###################
@@ -458,12 +444,10 @@ async def list_communities(
             limit=limit,
         )
 
-    # TODO: deprecate this
     @telemetry_event("get_communities")
     async def get_communities(
         self,
-        collection_id: Optional[UUID] = None,
-        levels: Optional[list[int]] = None,
+        collection_id: UUID,
         community_ids: Optional[list[int]] = None,
         offset: Optional[int] = None,
         limit: Optional[int] = None,
@@ -472,8 +456,8 @@ async def get_communities(
         return await self.providers.database.graph_handler.get_communities(
             collection_id=collection_id,
             community_ids=community_ids,
-            offset=offset or 0,
-            limit=limit or -1,
+            offset=offset,
+            limit=limit,
         )
 
     # @telemetry_event("create_new_graph")

py/core/providers/database/graph.py

@@ -544,7 +544,7 @@ async def get(
         entity_names: Optional[list[str]] = None,
         relationship_types: Optional[list[str]] = None,
         include_metadata: bool = False,
-    ) -> tuple[list[Relationship], int]:
+    ):
         """
         Get relationships from the specified store.
 
@@ -564,7 +564,7 @@ async def get(
         table_name = self._get_relationship_table_for_store(store_type)
 
         conditions = ["parent_id = $1"]
-        params = [parent_id]
+        params: list[Any] = [parent_id]
         param_index = 2
 
         if relationship_ids:

py/shared/abstractions/graph.py

@@ -90,9 +90,9 @@ class CommunityInfo(R2RSerializable):
     node: str
     cluster: UUID
     level: Optional[int]
-    id: Optional[UUID | int] = None
     parent_cluster: int | None
     is_final_cluster: bool
+    id: Optional[UUID | int] = None
     graph_id: Optional[UUID] = None
     collection_id: Optional[UUID] = None  # for backwards compatibility
     relationship_ids: Optional[list[UUID]] = None