Add KG tests

.github/workflows/integration-test-workflow-debian.yml

@@ -75,6 +75,22 @@ jobs:
         echo "Waiting for services to start..."
         sleep 30
 
+    - name: Run SDK GraphRAG
+      working-directory: ./py
+      run: |
+        poetry run python tests/integration/harness_sdk.py test_remove_all_files_and_ingest_sample_file_sdk
+        poetry run python tests/integration/harness_sdk.py test_kg_create_graph_sample_file_sdk
+        poetry run python tests/integration/harness_sdk.py test_kg_enrich_graph_sample_file_sdk
+        poetry run python tests/integration/harness_sdk.py test_kg_search_sample_file_sdk
+
+    - name: Run CLI GraphRAG
+      working-directory: ./py
+      run: |
+        poetry run python tests/integration/harness_sdk.py test_remove_all_files_and_ingest_sample_file_sdk
+        poetry run python tests/integration/harness_cli.py test_kg_create_graph_sample_file_cli
+        poetry run python tests/integration/harness_cli.py test_kg_enrichment_sample_file_cli
+        poetry run python tests/integration/harness_cli.py test_kg_search_sample_file_cli
+
     - name: Run CLI Ingestion
       working-directory: ./py
       run: |
@@ -117,8 +133,29 @@ jobs:
         poetry run python tests/integration/harness_sdk.py test_user_search_and_rag
         poetry run python tests/integration/harness_sdk.py test_user_password_management
         poetry run python tests/integration/harness_sdk.py test_user_profile_management
+        poetry run python tests/integration/harness_sdk.py test_user_overview
         poetry run python tests/integration/harness_sdk.py test_user_logout
 
+    - name: Run Collections
+      working-directory: ./py
+      run: |
+        poetry run python tests/integration/harness_sdk.py test_user_creates_collection
+        poetry run python tests/integration/harness_sdk.py test_user_updates_collection
+        poetry run python tests/integration/harness_sdk.py test_user_lists_collections
+        poetry run python tests/integration/harness_sdk.py test_user_collection_document_management
+        poetry run python tests/integration/harness_sdk.py test_user_removes_document_from_collection
+        poetry run python tests/integration/harness_sdk.py test_user_lists_documents_in_collection
+        poetry run python tests/integration/harness_sdk.py test_pagination_and_filtering
+        poetry run python tests/integration/harness_sdk.py test_advanced_collection_management
+        poetry run python tests/integration/harness_sdk.py test_user_gets_collection_details
+        poetry run python tests/integration/harness_sdk.py test_user_adds_user_to_collection
+        poetry run python tests/integration/harness_sdk.py test_user_removes_user_from_collection
+        poetry run python tests/integration/harness_sdk.py test_user_lists_users_in_collection
+        poetry run python tests/integration/harness_sdk.py test_user_gets_collections_for_user
+        poetry run python tests/integration/harness_sdk.py test_user_gets_collections_for_document
+        poetry run python tests/integration/harness_sdk.py test_user_permissions
+
+
     - name: Stop R2R server
       if: always()
       run: ps aux | grep "r2r serve" | awk '{print $2}' | xargs kill || true

py/cli/commands/ingestion.py

@@ -42,6 +42,7 @@ def ingest_files_from_urls(client, urls):
 
             files_to_ingest.append(temp_file.name)
             metadatas.append({"title": filename})
+            # TODO: use the utils function generate_document_id
             document_ids.append(uuid.uuid5(uuid.NAMESPACE_DNS, url))
 
         response = client.ingest_files(

py/cli/utils/docker_utils.py

@@ -243,7 +243,7 @@ def check_external_ollama(ollama_url="http://localhost:11434/api/version"):
 def check_set_docker_env_vars(exclude_postgres: bool = False):
 
     env_vars = {
-        "R2R_PROJECT_NAME": "r2r",
+        # "R2R_PROJECT_NAME": "r2r",
     }
 
     if not exclude_postgres:

py/core/base/providers/kg.py

@@ -81,6 +81,18 @@ async def add_kg_extractions(
         """Abstract method to add KG extractions."""
         pass
 
+    @abstractmethod
+    async def get_communities(
+        self,
+        collection_id: UUID,
+        offset: int,
+        limit: int,
+        levels: list[int] | None = None,
+        community_numbers: list[int] | None = None,
+    ) -> list[CommunityReport]:
+        """Abstract method to get communities."""
+        pass
+
     @abstractmethod
     async def get_entities(
         self,

py/core/main/api/kg_router.py

@@ -235,3 +235,35 @@ async def get_triples(
                 limit,
                 triple_ids,
             )
+
+        @self.router.get("/communities")
+        @self.base_endpoint
+        async def get_communities(
+            collection_id: UUID = Query(
+                ..., description="Collection ID to retrieve communities from."
+            ),
+            offset: int = Query(0, ge=0, description="Offset for pagination."),
+            limit: int = Query(
+                100, ge=1, le=1000, description="Limit for pagination."
+            ),
+            levels: Optional[list[int]] = Query(
+                None, description="Levels to filter by."
+            ),
+            community_numbers: Optional[list[int]] = Query(
+                None, description="Community numbers to filter by."
+            ),
+            auth_user=Depends(self.service.providers.auth.auth_wrapper),
+        ):
+            """
+            Retrieve communities from the knowledge graph.
+            """
+            if not auth_user.is_superuser:
+                logger.warning("Implement permission checks here.")
+
+            return await self.service.get_communities(
+                collection_id,
+                offset,
+                limit,
+                levels,
+                community_numbers,
+            )

py/core/main/api/management_router.py

@@ -370,8 +370,24 @@ async def document_chunks_app(
             is_owner = str(document_chunks_result[0].get("user_id")) == str(
                 auth_user.id
             )
+            document_collections = await self.service.document_collections(
+                document_uuid, 0, 1
+            )
 
-            if not is_owner and not auth_user.is_superuser:
+            user_has_access = (
+                is_owner
+                or set(auth_user.collection_ids).intersection(
+                    set(
+                        [
+                            ele.collection_id
+                            for ele in document_collections["results"]
+                        ]
+                    )
+                )
+                != set()
+            )
+
+            if not user_has_access and not auth_user.is_superuser:
                 raise R2RException(
                     "Only a superuser can arbitrarily call document_chunks.",
                     403,
@@ -512,7 +528,8 @@ async def list_collections_app(
         ) -> WrappedCollectionListResponse:
             if not auth_user.is_superuser:
                 raise R2RException(
-                    "Only a superuser can list all collections.", 403
+                    "Only a superuser can call the list collections endpoint.",
+                    403,
                 )
             list_collections_response = await self.service.list_collections(
                 offset=offset, limit=min(max(limit, 1), 1000)
@@ -543,7 +560,7 @@ async def add_user_to_collection_app(
             result = await self.service.add_user_to_collection(
                 user_uuid, collection_uuid
             )
-            return {"result": result}  # type: ignore
+            return result  # type: ignore
 
         @self.router.post("/remove_user_from_collection")
         @self.base_endpoint
@@ -611,7 +628,7 @@ async def get_collections_for_user_app(
             ),
             auth_user=Depends(self.service.providers.auth.auth_wrapper),
         ) -> WrappedUserCollectionResponse:
-            if str(auth_user.id) != user_id or not auth_user.is_superuser:
+            if str(auth_user.id) != user_id and not auth_user.is_superuser:
                 raise R2RException(
                     "The currently authenticated user does not have access to the specified collection.",
                     403,

py/core/main/orchestration/hatchet/ingestion_workflow.py

@@ -136,7 +136,7 @@ async def embed(self, context: Context) -> dict:
 
             return {
                 "status": "Successfully finalized ingestion",
-                "document_info": document_info.model_dump(),
+                "document_info": document_info.to_dict(),
             }
 
         @orchestration_provider.failure()

py/core/main/services/kg_service.py

@@ -294,3 +294,21 @@ async def get_triples(
             limit,
             triple_ids,
         )
+
+    @telemetry_event("get_communities")
+    async def get_communities(
+        self,
+        collection_id: UUID,
+        offset: int = 0,
+        limit: int = 100,
+        levels: Optional[list[int]] = None,
+        community_numbers: Optional[list[int]] = None,
+        **kwargs,
+    ):
+        return await self.providers.kg.get_communities(
+            collection_id,
+            offset,
+            limit,
+            levels,
+            community_numbers,
+        )

py/core/pipes/retrieval/kg_search_pipe.py

@@ -221,7 +221,7 @@ async def global_search(
         # map reduce
         async for message in input.message:
             map_responses = []
-            communities = self.kg_provider.get_communities(  # type: ignore
+            communities = await self.kg_provider.get_communities(  # type: ignore
                 level=kg_search_settings.kg_search_level
             )
 

py/core/providers/database/relational.py

@@ -1,6 +1,7 @@
+import asyncio
 import logging
 from contextlib import asynccontextmanager
-import asyncio
+
 import asyncpg
 
 from core.base import RelationalDBProvider

py/core/providers/kg/postgres.py

@@ -513,6 +513,34 @@ async def add_communities(self, communities: List[Any]) -> None:
             """
         await self.execute_many(QUERY, communities)
 
+    async def get_communities(
+        self,
+        collection_id: UUID,
+        offset: int = 0,
+        limit: int = 100,
+        levels: Optional[list[int]] = None,
+        community_numbers: Optional[list[int]] = None,
+    ) -> List[CommunityReport]:
+
+        query_parts = [
+            f"SELECT * FROM {self._get_table_name('community_report')} WHERE collection_id = $1 ORDER BY community_number LIMIT $2 OFFSET $3"
+        ]
+        params = [collection_id, limit, offset]
+
+        if levels is not None:
+            query_parts.append(f"AND level = ANY(${len(params) + 1})")
+            params.append(levels)
+
+        if community_numbers is not None:
+            query_parts.append(
+                f"AND community_number = ANY(${len(params) + 1})"
+            )
+            params.append(community_numbers)
+
+        QUERY = " ".join(query_parts)
+
+        return await self.fetch_query(QUERY, params)
+
     async def add_community_report(
         self, community_report: CommunityReport
     ) -> None:

py/pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "poetry.core.masonry.api"
 [tool.poetry]
 name = "r2r"
 readme = "README.md"
-version = "3.2.8"
+version = "3.2.9"
 
 description = "SciPhi R2R"
 authors = ["Owen Colegrove <owen@sciphi.ai>"]

py/sdk/auth.py

@@ -108,7 +108,6 @@ async def update_user(
             "bio": bio,
             "profile_picture": profile_picture,
         }
-        print("data = ", data)
         data = {k: v for k, v in data.items() if v is not None}
         return await client._make_request("PUT", "user", json=data)
 

py/sdk/kg.py

@@ -130,3 +130,34 @@ async def get_triples(
             params["triple_ids"] = ",".join(triple_ids)
 
         return await client._make_request("GET", "triples", params=params)
+
+    @staticmethod
+    async def get_communities(
+        client,
+        collection_id: str,
+        offset: int = 0,
+        limit: int = 100,
+        levels: Optional[list[int]] = None,
+        community_numbers: Optional[list[int]] = None,
+    ) -> dict:
+        """
+        Retrieve communities from the knowledge graph.
+
+        Args:
+            collection_id (str): The ID of the collection to retrieve communities from.
+
+        Returns:
+            dict: A dictionary containing the retrieved communities.
+        """
+        params = {
+            "collection_id": collection_id,
+            "offset": offset,
+            "limit": limit,
+        }
+
+        if levels:
+            params["levels"] = levels
+        if community_numbers:
+            params["community_numbers"] = community_numbers
+
+        return await client._make_request("GET", "communities", params=params)

py/shared/api/models/management/responses.py

@@ -151,10 +151,10 @@ class AddUserResponse(BaseModel):
 WrappedCollectionOverviewResponse = ResultsWrapper[
     list[CollectionOverviewResponse]
 ]
-WrappedAddUserResponse = ResultsWrapper[AddUserResponse]
+WrappedAddUserResponse = ResultsWrapper[None]
 WrappedUsersInCollectionResponse = PaginatedResultsWrapper[list[UserResponse]]
 WrappedUserCollectionResponse = PaginatedResultsWrapper[
-    list[CollectionOverviewResponse]
+    list[CollectionResponse]
 ]
 WrappedDocumentChunkResponse = PaginatedResultsWrapper[
     list[DocumentChunkResponse]