Patch/alternative up with nolan (#1642)

* prep for merge

* fixup crud

* rm pull

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

@@ -21,20 +21,20 @@
 from shared.api.models.kg.responses import (  # TODO: Need to review anything above this
     Community,
     Entity,
+    GraphResponse,
     KGEnrichmentResponse,
     KGTunePromptResponse,
     Relationship,
-    GraphResponse,
     WrappedCommunitiesResponse,
     WrappedCommunityResponse,
     WrappedEntitiesResponse,
     WrappedEntityResponse,
+    WrappedGraphResponse,
+    WrappedGraphsResponse,
     WrappedKGEnrichmentResponse,
     WrappedKGTunePromptResponse,
     WrappedRelationshipResponse,
     WrappedRelationshipsResponse,
-    WrappedGraphResponse,
-    WrappedGraphsResponse,
 )
 from shared.api.models.management.responses import (  # Document Responses; Prompt Responses; Chunk Responses; Conversation Responses; User Responses; TODO: anything below this hasn't been reviewed
     AnalyticsResponse,

py/core/main/api/v2/kg_router.py

@@ -11,7 +11,6 @@
 from core.base.api.models import (
     WrappedCommunitiesResponse,
     WrappedEntitiesResponse,
-    WrappedRelationshipsResponse,
     WrappedKGTunePromptResponse,
     WrappedRelationshipsResponse,
 )
@@ -66,7 +65,7 @@ def _register_workflows(self):
             )
         else:
             workflow_messages["extract-triples"] = (
-                "Document entities and relationships extracted successfully. To generate GraphRAG communities, run cluster on the collection this document belongs to."
+                "Document entities and relationships extracted successfully. To generate GraphRAG communities, POST to `/graphs/<collection_id>/communities/build` with a collection this document belongs to."
             )
             workflow_messages["build-communities"] = (
                 "Graph communities created successfully. You can view the communities at http://localhost:7272/v2/communities"

py/core/main/api/v3/documents_router.py

@@ -12,10 +12,7 @@
 from pydantic import Json
 
 from core.base import R2RException, RunType, generate_document_id
-from core.base.abstractions import (
-    KGCreationSettings,
-    KGRunType,
-)
+from core.base.abstractions import KGCreationSettings, KGRunType
 from core.base.api.models import (
     GenericBooleanResponse,
     WrappedBooleanResponse,
@@ -24,6 +21,7 @@
     WrappedDocumentResponse,
     WrappedDocumentsResponse,
     WrappedEntitiesResponse,
+    WrappedGenericMessageResponse,
     WrappedIngestionResponse,
     WrappedRelationshipsResponse,
 )
@@ -1262,7 +1260,7 @@ async def extract(
                 description="Whether to run the entities and relationships extraction process with orchestration.",
             ),
             auth_user=Depends(self.providers.auth.auth_wrapper),
-        ) -> WrappedIngestionResponse:
+        ) -> WrappedGenericMessageResponse:
             """
             Extracts entities and relationships from a document.
                 The entities and relationships extraction process involves:

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

@@ -6,9 +6,7 @@
 from fastapi import Body, Depends, Path, Query
 
 from core.base import R2RException, RunType
-from core.base.abstractions import (
-    KGRunType,
-)
+from core.base.abstractions import KGRunType
 from core.base.api.models import (
     GenericBooleanResponse,
     WrappedBooleanResponse,
@@ -450,7 +448,8 @@ async def reset(
                     403,
                 )
 
-            await self.services["kg"].delete_graph_v3(id=collection_id)
+            await self.services["kg"].reset_graph_v3(id=collection_id)
+            # await _pull(collection_id, auth_user)
             return GenericBooleanResponse(success=True)  # type: ignore
 
         # update graph
@@ -890,7 +889,7 @@ async def delete_entity(
                     403,
                 )
 
-            self.services["kg"].delete_entity(
+            await self.services["kg"].delete_entity(
                 parent_id=collection_id,
                 entity_id=entity_id,
             )
@@ -1444,7 +1443,7 @@ async def get_community(
             )
             print(f"results: {results}")
             if len(results) == 0 or len(results[0]) == 0:
-                raise R2RException("Relationship not found", 404)
+                raise R2RException("Community not found", 404)
             return results[0][0]
 
         @self.router.delete(

py/core/main/orchestration/simple/kg_workflow.py

@@ -91,6 +91,10 @@ async def extract_triples(input_data):
                     document_id=document_id,
                     **input_data["kg_creation_settings"],
                 ):
+                    print(
+                        "found extraction w/ entities = = ",
+                        len(extraction.entities),
+                    )
                     extractions.append(extraction)
                 await service.store_kg_extractions(extractions)
 
@@ -109,6 +113,16 @@ async def extract_triples(input_data):
     async def enrich_graph(input_data):
 
         input_data = get_input_data_dict(input_data)
+        workflow_status = await service.providers.database.get_workflow_status(
+            id=input_data.get("collection_id", None),
+            status_type="graph_cluster_status",
+        )
+        print("workflow_status = ", workflow_status)
+        if workflow_status == KGEnrichmentStatus.SUCCESS:
+            raise R2RException(
+                "Communities have already been built for this collection. To build communities again, first submit a POST request to `graphs/{collection_id}/reset`.",
+                400,
+            )
 
         try:
             num_communities = await service.kg_clustering(
@@ -142,11 +156,11 @@ async def enrich_graph(input_data):
                     input_data=input_data_copy,
                 )
 
-            # await service.providers.database.set_workflow_status(
-            #     id=input_data.get("collection_id", None),
-            #     status_type="graph_cluster_status",
-            #     status=KGEnrichmentStatus.SUCCESS,
-            # )
+            await service.providers.database.set_workflow_status(
+                id=input_data.get("collection_id", None),
+                status_type="graph_cluster_status",
+                status=KGEnrichmentStatus.SUCCESS,
+            )
             # return {
             #     "result": "successfully ran kg community summary workflows"
             # }

py/core/main/services/kg_service.py

@@ -24,6 +24,7 @@
     Graph,
     KGCreationSettings,
     KGEnrichmentSettings,
+    KGEnrichmentStatus,
     KGEntityDeduplicationSettings,
     KGEntityDeduplicationType,
     R2RException,
@@ -464,11 +465,28 @@ async def update_graph(
             description=description,
         )
 
+    @telemetry_event("reset_graph_v3")
+    async def reset_graph_v3(self, id: UUID) -> bool:
+        await self.providers.database.graph_handler.reset(
+            graph_id=id,
+        )
+        await self.providers.database.document_handler.set_workflow_status(
+            id=id,
+            status_type="graph_cluster_status",
+            status=KGEnrichmentStatus.PENDING,
+        )
+        return True
+
     @telemetry_event("delete_graph_v3")
     async def delete_graph_v3(self, id: UUID) -> bool:
         await self.providers.database.graph_handler.delete(
             graph_id=id,
         )
+        # await self.providers.database.document_handler.set_workflow_status(
+        #     id=id,
+        #     status_type="graph_cluster_status",
+        #     status=KGEnrichmentStatus.PENDING,
+        # )
         return True
 
     @telemetry_event("get_document_ids_for_create_graph")

py/core/pipes/retrieval/kg_search_pipe.py

@@ -179,6 +179,14 @@ async def search(
                     # "document_ids",
                 ],
             ):
+                try:
+                    # TODO - remove this nasty hack
+                    search_result["metadata"] = json.loads(
+                        search_result["metadata"]
+                    )
+                except:
+                    pass
+
                 yield GraphSearchResult(
                     content=KGRelationshipResult(
                         # name=search_result["name"],

py/core/providers/database/collection.py

@@ -46,7 +46,7 @@ async def create_tables(self) -> None:
             user_id UUID,
             name TEXT NOT NULL,
             description TEXT,
-            graph_cluster_status TEXT DEFAULT 'PENDING',
+            graph_cluster_status TEXT DEFAULT 'pending',
             created_at TIMESTAMPTZ DEFAULT NOW(),
             updated_at TIMESTAMPTZ DEFAULT NOW()
         );