From ed7f7344e31d84da92064fda32f750a7db68ec04 Mon Sep 17 00:00:00 2001 From: NolanTrem <34580718+NolanTrem@users.noreply.github.com> Date: Fri, 29 Nov 2024 22:28:00 -0600 Subject: [PATCH] Clean up type errors, cleaner code --- js/sdk/src/v3/clients/graphs.ts | 25 +++++++++++ py/core/main/api/v3/graph_router.py | 68 ++++++++++++++++++++++++----- py/core/main/services/kg_service.py | 46 +++++++------------ py/core/providers/database/graph.py | 4 +- py/shared/abstractions/graph.py | 2 +- 5 files changed, 101 insertions(+), 44 deletions(-) diff --git a/js/sdk/src/v3/clients/graphs.ts b/js/sdk/src/v3/clients/graphs.ts index 85734c554..351b0b2e3 100644 --- a/js/sdk/src/v3/clients/graphs.ts +++ b/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; + runWithOrchestration?: boolean; }): Promise { return this.client.makeRequest( "POST", diff --git a/py/core/main/api/v3/graph_router.py b/py/core/main/api/v3/graph_router.py index eeb695fae..a8f189d23 100644 --- a/py/core/main/api/v3/graph_router.py +++ b/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 diff --git a/py/core/main/services/kg_service.py b/py/core/main/services/kg_service.py index f29553f32..b5c446f66 100644 --- a/py/core/main/services/kg_service.py +++ b/py/core/main/services/kg_service.py @@ -207,12 +207,11 @@ 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, @@ -220,19 +219,6 @@ async def delete_entity_v3( 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") diff --git a/py/core/providers/database/graph.py b/py/core/providers/database/graph.py index 6ceefe7b1..be559217a 100644 --- a/py/core/providers/database/graph.py +++ b/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: diff --git a/py/shared/abstractions/graph.py b/py/shared/abstractions/graph.py index d7fe6548b..71221546f 100644 --- a/py/shared/abstractions/graph.py +++ b/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